A high-performance general-purpose compute library
graphics/gravity_sim.cpp
/*******************************************************
* Copyright (c) 2014, ArrayFire
* All rights reserved.
*
* This file is distributed under 3-clause BSD license.
* The complete license agreement can be obtained at:
* http://arrayfire.com/licenses/BSD-3-Clause
********************************************************/
#include <arrayfire.h>
#include <cstdio>
#include <iostream>
#include <vector>
#include "gravity_sim_init.h"
using namespace af;
using namespace std;
static const bool is3D = true;
const static int total_particles = 4000;
static const int reset = 3000;
static const float min_dist = 3;
static const int width = 768, height = 768, depth = 768;
static const int gravity_constant = 20000;
float mass_range = 0;
float min_mass = 0;
void initial_conditions_rand(af::array &mass, vector<af::array> &pos,
vector<af::array> &vels,
vector<af::array> &forces) {
for (int i = 0; i < (int)pos.size(); ++i) {
pos[i] = af::randn(total_particles) * width + width;
vels[i] = 0 * af::randu(total_particles) - 0.5;
forces[i] = af::constant(0, total_particles);
}
mass = af::constant(gravity_constant, total_particles);
}
void initial_conditions_galaxy(af::array &mass, vector<af::array> &pos,
vector<af::array> &vels,
vector<af::array> &forces) {
af::array initial_cond_consts(af::dim4(7, total_particles), hbd);
initial_cond_consts = initial_cond_consts.T();
for (int i = 0; i < (int)pos.size(); ++i) {
pos[i] = af::randn(total_particles) * width + width;
vels[i] = 0 * (af::randu(total_particles) - 0.5);
forces[i] = af::constant(0, total_particles);
}
mass = initial_cond_consts(span, 0);
pos[0] = (initial_cond_consts(span, 1) / 32 + 0.6) * width;
pos[1] = (initial_cond_consts(span, 2) / 32 + 0.3) * height;
pos[2] = (initial_cond_consts(span, 3) / 32 + 0.5) * depth;
vels[0] = (initial_cond_consts(span, 4) / 32) * width;
vels[1] = (initial_cond_consts(span, 5) / 32) * height;
vels[2] = (initial_cond_consts(span, 6) / 32) * depth;
pos[0](seq(0, pos[0].dims(0) - 1, 2)) -= 0.4 * width;
pos[1](seq(0, pos[0].dims(0) - 1, 2)) += 0.4 * height;
vels[0](seq(0, pos[0].dims(0) - 1, 2)) += 4;
min_mass = min<float>(mass);
mass_range = max<float>(mass) - min<float>(mass);
}
af::array ids_from_pos(vector<af::array> &pos) {
return (pos[0].as(u32) * height) + pos[1].as(u32);
}
af::array ids_from_3D(vector<af::array> &pos, float Rx, float Ry, float Rz) {
af::array x0 = (pos[0] - width / 2);
af::array y0 =
(pos[1] - height / 2) * cos(Rx) + (pos[2] - depth / 2) * sin(Rx);
af::array z0 =
(pos[2] - depth / 2) * cos(Rx) - (pos[2] - depth / 2) * sin(Rx);
af::array x1 = x0 * cos(Ry) - z0 * sin(Ry);
af::array y1 = y0;
af::array x2 = x1 * cos(Rz) + y1 * sin(Rz);
af::array y2 = y1 * cos(Rz) - x1 * sin(Rz);
x2 += width / 2;
y2 += height / 2;
return (x2.as(u32) * height) + y2.as(u32);
}
af::array ids_from_3D(vector<af::array> &pos, float Rx, float Ry, float Rz,
af::array filter) {
af::array x0 = (pos[0](filter) - width / 2);
af::array y0 = (pos[1](filter) - height / 2) * cos(Rx) +
(pos[2](filter) - depth / 2) * sin(Rx);
af::array z0 = (pos[2](filter) - depth / 2) * cos(Rx) -
(pos[2](filter) - depth / 2) * sin(Rx);
af::array x1 = x0 * cos(Ry) - z0 * sin(Ry);
af::array y1 = y0;
af::array x2 = x1 * cos(Rz) + y1 * sin(Rz);
af::array y2 = y1 * cos(Rz) - x1 * sin(Rz);
x2 += width / 2;
y2 += height / 2;
return (x2.as(u32) * height) + y2.as(u32);
}
void simulate(af::array &mass, vector<af::array> &pos, vector<af::array> &vels,
vector<af::array> &forces, float dt) {
for (int i = 0; i < (int)pos.size(); ++i) {
pos[i] += vels[i] * dt;
pos[i].eval();
}
// calculate forces to each particle
vector<af::array> diff(pos.size());
af::array dist = af::constant(0, pos[0].dims(0), pos[0].dims(0));
for (int i = 0; i < (int)pos.size(); ++i) {
diff[i] = tile(pos[i], 1, pos[i].dims(0)) -
transpose(tile(pos[i], 1, pos[i].dims(0)));
dist += (diff[i] * diff[i]);
}
dist = sqrt(dist);
dist = af::max(min_dist, dist);
dist *= dist * dist;
for (int i = 0; i < (int)pos.size(); ++i) {
// calculate force vectors
forces[i] = diff[i] / dist;
forces[i].eval();
// af::array idx = af::where(af::isNaN(forces[i]));
// if(idx.elements() > 0)
// forces[i](idx) = 0;
// forces[i] = sum(forces[i]).T();
forces[i] = matmul(forces[i].T(), mass);
// update force scaled to time, magnitude constant
forces[i] *= (gravity_constant);
forces[i].eval();
// update velocities from forces
vels[i] += forces[i] * dt;
vels[i].eval();
// noise
// forces[i] += 0.1 * af::randn(forces[i].dims(0));
// dampening
// vels[i] *= 1 - (0.005*dt);
}
}
void collisions(vector<af::array> &pos, vector<af::array> &vels, bool is3D) {
// clamp particles inside screen border
af::array invalid_x = -2 * (pos[0] > width - 1 || pos[0] < 0) + 1;
af::array invalid_y = -2 * (pos[1] > height - 1 || pos[1] < 0) + 1;
// af::array invalid_x = (pos[0] < width-1 || pos[0] > 0);
// af::array invalid_y = (pos[1] < height-1 || pos[1] > 0);
vels[0] = invalid_x * vels[0];
vels[1] = invalid_y * vels[1];
af::array projected_px = min(width - 1, max(0, pos[0]));
af::array projected_py = min(height - 1, max(0, pos[1]));
pos[0] = projected_px;
pos[1] = projected_py;
if (is3D) {
af::array invalid_z = -2 * (pos[2] > depth - 1 || pos[2] < 0) + 1;
vels[2] = invalid_z * vels[2];
af::array projected_pz = min(depth - 1, max(0, pos[2]));
pos[2] = projected_pz;
}
}
int main(int, char **) {
try {
af::Window myWindow(width, height,
"Gravity Simulation using ArrayFire");
myWindow.setColorMap(AF_COLORMAP_HEAT);
int frame_count = 0;
// Initialize the kernel array just once
const af::array draw_kernel = gaussianKernel(7, 7);
const int dims = (is3D) ? 3 : 2;
vector<af::array> pos(dims);
vector<af::array> vels(dims);
vector<af::array> forces(dims);
af::array mass;
// Generate a random starting state
initial_conditions_galaxy(mass, pos, vels, forces);
af::array image = af::constant(0, width, height);
af::array ids(total_particles, u32);
while (!myWindow.close()) {
float dt = af::timer::stop(timer);
af::array mid = mass(span) > (min_mass + mass_range / 3);
ids = (is3D) ? ids_from_3D(pos, 0, 0 + frame_count / 150.f, 0, mid)
: ids_from_pos(pos);
// ids = (is3D)? ids_from_3D(pos, 0, 0, 0, mid) : ids_from_pos(pos);
// //uncomment for no 3d rotation
image(ids) += 4.f;
mid = mass(span) > (min_mass + 2 * mass_range / 3);
ids = (is3D) ? ids_from_3D(pos, 0, 0 + frame_count / 150.f, 0, mid)
: ids_from_pos(pos);
// ids = (is3D)? ids_from_3D(pos, 0, 0, 0, mid) : ids_from_pos(pos);
// //uncomment for no 3d rotation
image(ids) += 4.f;
ids = (is3D) ? ids_from_3D(pos, 0, 0 + frame_count / 150.f, 0)
: ids_from_pos(pos);
// ids = (is3D)? ids_from_3D(pos, 0, 0, 0) : ids_from_pos(pos);
// //uncomment for no 3d rotation
image(ids) += 4.f;
image = convolve(image, draw_kernel);
myWindow.image(image);
image = af::constant(0, image.dims());
frame_count++;
// Generate a random starting state
if (frame_count % reset == 0) {
initial_conditions_galaxy(mass, pos, vels, forces);
}
// simulate
simulate(mass, pos, vels, forces, dt);
// check for collisions and adjust positions/velocities accordingly
collisions(pos, vels, is3D);
}
} catch (af::exception &e) {
fprintf(stderr, "%s\n", e.what());
throw;
}
return 0;
}
Window object to render af::arrays.
Definition: graphics.h:37
A multi dimensional data container.
Definition: array.h:37
dim4 dims() const
Get dimensions of the array.
const array as(dtype type) const
Casts the array into another data type.
void eval() const
Evaluate any JIT expressions to generate data for the array.
Generic object that represents size and shape.
Definition: dim4.hpp:26
An ArrayFire exception class.
Definition: exception.h:22
virtual const char * what() const
Returns an error message for the exception in a string format.
Definition: exception.h:46
seq is used to create sequences for indexing af::array
Definition: seq.h:46
@ u32
32-bit unsigned integral values
Definition: defines.h:217
@ AF_COLORMAP_HEAT
Heat map.
Definition: defines.h:460
AFAPI array cos(const array &in)
C++ Interface to evaluate the cosine function.
AFAPI array sin(const array &in)
C++ Interface to evaluate the sine function.
AFAPI array sqrt(const array &in)
C++ Interface to evaluate the square root.
AFAPI array matmul(const array &lhs, const array &rhs, const matProp optLhs=AF_MAT_NONE, const matProp optRhs=AF_MAT_NONE)
C++ Interface to multiply two matrices.
AFAPI array transpose(const array &in, const bool conjugate=false)
C++ Interface to transpose a matrix.
array constant(T val, const dim4 &dims, const dtype ty=(af_dtype) dtype_traits< T >::ctype)
C++ Interface to generate an array with elements set to a specified value.
AFAPI void info()
AFAPI array gaussianKernel(const int rows, const int cols, const double sig_r=0, const double sig_c=0)
C++ Interface for generating gausian kernels.
AFAPI array tile(const array &in, const unsigned x, const unsigned y=1, const unsigned z=1, const unsigned w=1)
C++ Interface to generate a tiled array.
AFAPI array randn(const dim4 &dims, const dtype ty, randomEngine &r)
C++ Interface to create an array of random numbers normally distributed.
AFAPI array randu(const dim4 &dims, const dtype ty, randomEngine &r)
C++ Interface to create an array of random numbers uniformly distributed.
AFAPI array max(const array &in, const int dim=-1)
C++ Interface to return the maximum along a given dimension.
AFAPI array min(const array &in, const int dim=-1)
C++ Interface to return the minimum along a given dimension.
AFAPI array convolve(const array &signal, const array &filter, const convMode mode=AF_CONV_DEFAULT, const convDomain domain=AF_CONV_AUTO)
C++ Interface for convolution any(one through three) dimensional signals.
Definition: algorithm.h:15
AFAPI array filter(const array &image, const array &kernel)
Image Filtering.
Internal timer object.
Definition: timing.h:29
static AFAPI timer start()
static AFAPI double stop()