def test_2d(self):
data = np.arange(6, dtype=np.float64).reshape((3, 2))
kernel = cgen.Assign("output_grid[i2][i1]", "input_grid[i2][i1] + 3")
propagator = Propagator("process", 3, (2, ), [])
propagator.add_param("input_grid", data.shape, data.dtype)
propagator.add_param("output_grid", data.shape, data.dtype)
propagator.loop_body = kernel
f = propagator.cfunction
arr = np.empty_like(data)
f(data, arr)
assert (arr[2][1] == 8)
def test_4d(self):
kernel = cgen.Assign("output_grid[i4][i1][i2][i3]",
"input_grid[i4][i1][i2][i3] + 3")
data = np.arange(120, dtype=np.float64).reshape((5, 4, 3, 2))
propagator = Propagator("process", 5, (4, 3, 2), [])
propagator.add_param("input_grid", data.shape, data.dtype)
propagator.add_param("output_grid", data.shape, data.dtype)
propagator.loop_body = kernel
f = propagator.cfunction
arr = np.empty_like(data)
f(data, arr)
assert (arr[4][3][2][1] == 122)
def test_value_param():
data = np.arange(6, dtype=np.float64).reshape((3, 2))
kernel = Assign("output_grid[i2][i1]", "input_grid[i2][i1] + offset")
propagator = Propagator("process", 3, (2, ), [])
propagator.add_param("input_grid", data.shape, data.dtype)
propagator.add_param("output_grid", data.shape, data.dtype)
propagator.add_scalar_param("offset", np.int32)
propagator.loop_body = kernel
f = propagator.cfunction
arr = np.empty_like(data)
f(data, arr, np.int32(3))
assert (arr[2][1] == 8)