def test_subslice_shape_i(self):
"""Subslice shape reduced by one dimension (two remain)"""
ND, S0, S1 = range(3)
@function(x=self.DataSlice, i=Long, v=Slice(Long))
def get_i(x, i, v):
s = x[i]
v[ND] = s.ndim
v[S0] = s.shape[0]
v[S1] = s.shape[1]
m = module([get_i])
v = (Long.c_type * 3)()
# Shape and dimensions should not depend on indices.
for i in range(5):
m.get_i(self.data, i, v)
self.assertEqual(v[ND], 2)
self.assertEqual(v[S0], 2)
self.assertEqual(v[S1], 3)
def test_reference_arg(self):
"""Slice is treated as reference type."""
from nitrous.types import is_aggregate
self.assertTrue(is_aggregate(Slice(Long)))
import unittest
from nitrous.module import module
from nitrous.function import function
from nitrous.types import Double, Long, Structure
from nitrous.types.array import Slice, Any
Coord = Structure("Coord", ("x", Double), ("y", Double), ("z", Double))
Coords = Slice(Coord)
class StructureTests(unittest.TestCase):
def test_repr(self):
self.assertEqual(repr(Coord), "<Structure 'Coord', 3 fields>")
def test_load_fields(self):
import ctypes
# Attribute load from subscript
@function(Double, a=Coords, i=Long)
def sum_1(a, i):
return a[i].x + a[i].y + a[i].z
# Subscript, reference assignment, then attribute load
@function(Double, a=Coords, i=Long)
def sum_2(a, i):
ai = a[i]
return ai.x + ai.y + ai.z
m = module([sum_1, sum_2])
a = (Coord.c_type * 2)((1, 2, 3), (4, 5, 6))
import unittest
from nitrous.module import module
from nitrous.function import function
from nitrous.types import Float
from nitrous.types.array import Slice
from nitrous.exp.vector import Vector, load, store, get_element, set_element, fill
FloatP = Slice(Float, (4, ))
Float4 = Vector(Float, 4)
load4f = load(Float4)
store4f = store(Float4)
get4f = get_element(Float4)
set4f = set_element(Float4)
fill4f = fill(Float4)
@function(Float, a=Float, b=Float, c=Float, d=Float)
def hadd4(a, b, c, d):
v = Float4()
v = set4f(v, 0, a)
v = set4f(v, 1, b)
v = set4f(v, 2, c)
v = set4f(v, 3, d)
return get4f(v, 0) + get4f(v, 1) + get4f(v, 2) + get4f(v, 3)
import unittest
import numpy as np
from nitrous.function import function
from nitrous.types import Double, Index
from nitrous.types.array import Array, FastSlice, Slice, Any
DoubleNx3 = Slice(Double, shape=(Any, 3))
DoubleN = Slice(Double)
DoubleNArray = FastSlice(Double)
Double3 = Array(Double, (3, ))
X, Y, Z = range(3)
@function(d=DoubleNx3, out=DoubleNArray, n=Index, m=Index)
def sum_1(d, out, n, m):
for k in range(m):
for i in range(n):
for j in range(i + 1, n):
out[X] += d[i, X] * d[j, X]
out[Y] += d[i, Y] * d[j, Y]
out[Z] += d[i, Z] * d[j, Z]
@function(d=DoubleNx3, out=DoubleN, n=Index, m=Index)
def sum_2(d, out, n, m):
for k in range(m):
def test(self):
"""N-body benchmark adaptation from http://shootout.alioth.debian.org"""
from nitrous.module import module
from nitrous.function import function
from nitrous.types import Long, Double
from nitrous.types.array import Slice, Any
from nitrous.lib.math import sqrt
X, Y, Z = range(3)
PI = 3.14159265358979323
SOLAR_MASS = 4 * PI * PI
DAYS_PER_YEAR = 365.24
# In the following sequence
# - Sun
# - Jupiter
# - Saturn
# - Uranus
# - Neptune
DoubleNx3 = Slice(Double, shape=(Any, 3))
DoubleN = Slice(Double)
common_args = {
"xyz": DoubleNx3,
"vxyz": DoubleNx3,
"mass": DoubleN,
"n_bodies": Long
}
@function(vxyz=DoubleNx3, mass=DoubleN, n_bodies=Long)
def offset_momentum(vxyz, mass, n_bodies):
px = 0.0
py = 0.0
pz = 0.0
for i in range(n_bodies):
px -= vxyz[i, X] * mass[i]
py -= vxyz[i, Y] * mass[i]
pz -= vxyz[i, Z] * mass[i]
vxyz[0, X] = px / SOLAR_MASS
vxyz[0, Y] = py / SOLAR_MASS
vxyz[0, Z] = pz / SOLAR_MASS
@function(Double, **common_args)
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
@function(dt=Double, **common_args)
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]
@function(n_steps=Long, **common_args)
def loop(xyz, vxyz, mass, n_bodies, n_steps):
for i in range(n_steps):
advance(xyz, vxyz, mass, n_bodies, 0.01)
m = module([offset_momentum, energy, loop])
xyz = np.genfromtxt("tests/data/nbody-position")
vxyz = np.genfromtxt("tests/data/nbody-velocity") * DAYS_PER_YEAR
mass = np.genfromtxt("tests/data/nbody-mass") * SOLAR_MASS
m.offset_momentum(vxyz, mass, 5)
# from time import time
# t0 = time()
e0 = m.energy(xyz, vxyz, mass, 5)
self.assertAlmostEqual(e0, -0.169075164)
# print " e=", e0, "Elapsed", time() - t0
m.loop(xyz, vxyz, mass, 5, 50000000)
e1 = m.energy(xyz, vxyz, mass, 5)
self.assertAlmostEqual(e1, -0.169059907)