Cwerg/FE/TestData/nbody_test.cw at master · robertmuth/Cwerg · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
; adapted from http://shootout.alioth.debian.org/
module:

import test

import cmp

import fmt

import fmt_real

global PI = 3.141592653589793_r64

global SOLAR_MASS = PI * PI * 4.0

global DAYS_PER_YEAR = 365.24_r64

global REL_ERR = 1e-13_r64

rec Body:
    x r64
    y r64
    z r64
    vx r64
    vy r64
    vz r64
    mass r64

fun BodyCommon(x r64, y r64, z r64, vx r64, vy r64, vz r64, m r64) Body:
    return {Body: x, y, z, vx * DAYS_PER_YEAR, vy * DAYS_PER_YEAR,
            vz * DAYS_PER_YEAR, m * SOLAR_MASS}

fun BodyJupiter() Body:
    return BodyCommon(4.84143144246472090e+00, -1.16032004402742839e+00,
             -1.03622044471123109e-01, 1.66007664274403694e-03,
             7.69901118419740425e-03, -6.90460016972063023e-05,
             9.54791938424326609e-04)

fun BodySaturn() Body:
    return BodyCommon(8.34336671824457987e+00, 4.12479856412430479e+00,
             -4.03523417114321381e-01, -2.76742510726862411e-03,
             4.99852801234917238e-03, 2.30417297573763929e-05,
             2.85885980666130812e-04)

fun BodyUranus() Body:
    return BodyCommon(1.28943695621391310e+01, -1.51111514016986312e+01,
             -2.23307578892655734e-01, 2.96460137564761618e-03,
             2.37847173959480950e-03, -2.96589568540237556e-05,
             4.36624404335156298e-05)

fun BodyNeptune() Body:
    return BodyCommon(1.53796971148509165e+01, -2.59193146099879641e+01,
             1.79258772950371181e-01, 2.68067772490389322e-03,
             1.62824170038242295e-03, -9.51592254519715870e-05,
             5.15138902046611451e-05)

fun BodySun() Body:
    return BodyCommon(0, 0, 0, 0, 0, 0, 1.0)

fun UpdateOffsetMomentum(bodies ^![5]Body) void:
    let! px = 0.0_r64
    let! py = 0.0_r64
    let! pz = 0.0_r64
    for i = 0, len(bodies^), 1:
        let b = @!bodies^[i]
        set px += b^.vx * b^.mass
        set py += b^.vy * b^.mass
        set pz += b^.vz * b^.mass
    let s = @!bodies^[0]
    set s^.vx = -(px / SOLAR_MASS)
    set s^.vy = -(py / SOLAR_MASS)
    set s^.vz = -(pz / SOLAR_MASS)

fun Advance(bodies ^![5]Body, dt r64) void:
    for i = 0, len(bodies^), 1:
        let bi = @!bodies^[i]
        for j = i + 1, len(bodies^), 1:
            let bj = @!bodies^[j]
            let dx = bi^.x - bj^.x
            let dy = bi^.y - bj^.y
            let dz = bi^.z - bj^.z
            let d2 = dx * dx + dy * dy + dz * dz
            let d = sqrt(d2)
            let mag = dt / (d * d2)
            let mj = bj^.mass * mag
            set bi^.vx -= dx * mj
            set bi^.vy -= dy * mj
            set bi^.vz -= dz * mj
            let mi = bi^.mass * mag
            set bj^.vx += dx * mi
            set bj^.vy += dy * mi
            set bj^.vz += dz * mi
    for i = 0, len(bodies^), 1:
        let bi = @!bodies^[i]
        set bi^.x += dt * bi^.vx
        set bi^.y += dt * bi^.vy
        set bi^.z += dt * bi^.vz

fun Energy(bodies ^[5]Body) r64:
    let! e = 0.0_r64
    for i = 0, len(bodies^), 1:
        let bi = @bodies^[i]
        set e += 0.5 * bi^.mass *
          (bi^.vx * bi^.vx + bi^.vy * bi^.vy + bi^.vz * bi^.vz)
        for j = i + 1, len(bodies^), 1:
            let bj = @bodies^[j]
            let dx = bi^.x - bj^.x
            let dy = bi^.y - bj^.y
            let dz = bi^.z - bj^.z
            let d = sqrt(dx * dx + dy * dy + dz * dz)
            set e -= bi^.mass * bj^.mass / d
    return e

global DT = 0.01_r64

global NUM_ITER = 250000_u32

fun main(argc s32, argv ^^u8) s32:
    ref let! bodies = {[5]Body: BodySun(), BodyJupiter(), BodySaturn(),
                       BodyUranus(), BodyNeptune()}
    do UpdateOffsetMomentum(@!bodies)
    ; do permute(@!v, DIM)
    ; DIM! = 5040
    ; test\AssertEq#(COUNT, 5040_u32)
    if true:
        ; sanity test with one iteration
        do Advance(@!bodies, DT)
        let e = Energy(@bodies)
        ; fmt\print#(wrap_as(e, fmt\r64_hex), " ", e, "\n")
        test\AssertGenericEq#({cmp\r64r: -0.16907495402506745, REL_ERR},
                               {cmp\r64r: e})
    else:
        for i = 0, NUM_ITER, 1:
            do Advance(@!bodies, DT)
        let e = Energy(@bodies)
        ; fmt\print#(wrap_as(e, fmt\r64_hex), " ", e,  "\n")
        test\AssertGenericEq#({cmp\r64r: -0.1690859889909308, REL_ERR},
                               {cmp\r64r: e})
    test\Success#()
    return 0