def test_sqrt(self): @function(Double, x=Double) def sqrt(x): return nitrous.lib.math.sqrt(Double)(x) m = module([sqrt]) self.assertAlmostEqual(math.sqrt(10.0), m.sqrt(10.0))
def energy(xyz, vxyz, mass, n_bodies): e = 0.0 for i in range(n_bodies): vx = vxyz[i, X] vy = vxyz[i, Y] vz = vxyz[i, Z] e += 0.5 * mass[i] * (vx * vx + vy * vy + vz * vz) for j in range(i + 1, n_bodies): dx = xyz[i, X] - xyz[j, X] dy = xyz[i, Y] - xyz[j, Y] dz = xyz[i, Z] - xyz[j, Z] d2 = dx * dx + dy * dy + dz * dz e -= mass[i] * mass[j] / sqrt(Double)(d2) return e
def advance(xyz, vxyz, mass, n_bodies, dt): for i in range(n_bodies): for j in range(i + 1, n_bodies): dx = xyz[i, X] - xyz[j, X] dy = xyz[i, Y] - xyz[j, Y] dz = xyz[i, Z] - xyz[j, Z] d2 = dx * dx + dy * dy + dz * dz mag = dt / (d2 * sqrt(Double)(d2)) vxyz[i, X] -= dx * mass[j] * mag vxyz[i, Y] -= dy * mass[j] * mag vxyz[i, Z] -= dz * mass[j] * mag vxyz[j, X] += dx * mass[i] * mag vxyz[j, Y] += dy * mass[i] * mag vxyz[j, Z] += dz * mass[i] * mag for i in range(n_bodies): xyz[i, X] += dt * vxyz[i, X] xyz[i, Y] += dt * vxyz[i, Y] xyz[i, Z] += dt * vxyz[i, Z]
def _norm(a): return sqrt(Double)(a.x * a.x + a.y * a.y + a.z * a.z)