Skip to content

Instantly share code, notes, and snippets.

@iffy
Created October 26, 2022 17:49
Show Gist options
  • Save iffy/04d5f1246697b5dcac0539ae8ec2c8f3 to your computer and use it in GitHub Desktop.
Save iffy/04d5f1246697b5dcac0539ae8ec2c8f3 to your computer and use it in GitHub Desktop.
I wanted to see how Nim compared (readability and speedwise) with what's reported here: https://jott.live/markdown/py3.11_vs_3.8
## Compile for speed with `nim c -d:danger nbody.nim`
## Run `./nbody 10000000`
import std/os
import std/sequtils
import std/strformat
import std/strutils
import std/tables
from std/math import pow
type
Vec3 = array[3, float]
BodyDef = ref object
pos: Vec3
vel: Vec3
mass: float
Pair = ref object
a: BodyDef
b: BodyDef
const
PI = 3.14159265358979323
SOLAR_MASS = 4 * PI * PI
DAYS_PER_YEAR = 365.24
template `x`(v: Vec3): float = v[0]
template `y`(v: Vec3): float = v[1]
template `z`(v: Vec3): float = v[2]
func `+`(a, b: Vec3): Vec3 = [a.x + b.x, a.y + b.y, a.z + b.z]
func `-`(a, b: Vec3): Vec3 = [a.x - b.x, a.y - b.y, a.z - b.z]
func `*`(a: Vec3, s: float): Vec3 = [a.x * s, a.y * s, a.z * s]
func `/`(a: Vec3, s: float): Vec3 = [a.x / s, a.y / s, a.z / s]
func combinations(l: seq[BodyDef]): seq[Pair] =
for x in 0..<(l.len - 1):
let ls = l[x+1..^1]
for y in ls:
result.add(Pair(a: l[x], b: y))
proc toBodyDef(pos: Vec3, vel: Vec3, mass: float): BodyDef =
new(result)
result.pos = pos
result.vel = vel
result.mass = mass
var
BODIES: Table[string, BodyDef] = {
"sun": toBodyDef([0.0, 0.0, 0.0], [0.0, 0.0, 0.0], SOLAR_MASS),
"jupiter": toBodyDef([4.84143144246472090e+00,
-1.16032004402742839e+00,
-1.03622044471123109e-01],
[1.66007664274403694e-03 * DAYS_PER_YEAR,
7.69901118419740425e-03 * DAYS_PER_YEAR,
-6.90460016972063023e-05 * DAYS_PER_YEAR],
9.54791938424326609e-04 * SOLAR_MASS),
"saturn": toBodyDef([8.34336671824457987e+00,
4.12479856412430479e+00,
-4.03523417114321381e-01],
[-2.76742510726862411e-03 * DAYS_PER_YEAR,
4.99852801234917238e-03 * DAYS_PER_YEAR,
2.30417297573763929e-05 * DAYS_PER_YEAR],
2.85885980666130812e-04 * SOLAR_MASS),
"uranus": toBodyDef([1.28943695621391310e+01,
-1.51111514016986312e+01,
-2.23307578892655734e-01],
[2.96460137564761618e-03 * DAYS_PER_YEAR,
2.37847173959480950e-03 * DAYS_PER_YEAR,
-2.96589568540237556e-05 * DAYS_PER_YEAR],
4.36624404335156298e-05 * SOLAR_MASS),
"neptune": toBodyDef([1.53796971148509165e+01,
-2.59193146099879641e+01,
1.79258772950371181e-01],
[2.68067772490389322e-03 * DAYS_PER_YEAR,
1.62824170038242295e-03 * DAYS_PER_YEAR,
-9.51592254519715870e-05 * DAYS_PER_YEAR],
5.15138902046611451e-05 * SOLAR_MASS)
}.toTable
SYSTEM = toSeq(BODIES.values())
PAIRS = combinations(SYSTEM)
proc advance(dt: float, n: Natural, bodies: var seq[BodyDef], pairs: var seq[Pair]) =
for i in 0..<n:
for pair in pairs.mitems:
let
a = pair.a
b = pair.b
d = a.pos - b.pos
mag = dt * pow(d.x * d.x + d.y * d.y + d.z * d.z, -1.5)
b1m = a.mass * mag
b2m = b.mass * mag
pair.a.vel = pair.a.vel - d * b2m
pair.b.vel = pair.b.vel + d * b1m
for body in bodies.mitems:
body.pos = body.pos + (body.vel * dt)
proc report_energy(bodies: var seq[BodyDef], pairs: var seq[Pair], e=0.0) =
var e = e
for pair in pairs.mitems:
let
a = pair.a
b = pair.b
d = a.pos - b.pos
e -= (a.mass * b.mass) / pow(d.x * d.x + d.y * d.y + d.z * d.z, 0.5)
for body in bodies:
let v = body.vel
e += body.mass * (v.x * v.x + v.y * v.y + v.z * v.z) / 2.0
echo &"{e:.9f}"
proc offset_momentum(refpoint: var BodyDef, bodies: var seq[BodyDef]) =
var p: Vec3 = [0.0, 0.0, 0.0]
for body in bodies.mitems:
p = p - body.vel * body.mass
refpoint.vel = p / refpoint.mass
proc main(n: Natural, refname="sun") =
offset_momentum(BODIES[refname], SYSTEM)
report_energy(SYSTEM, PAIRS)
advance(0.01, n, SYSTEM, PAIRS)
report_energy(SYSTEM, PAIRS)
when isMainModule:
main(paramStr(1).parseInt)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment