def test_vectorization_fixed_size():
configurations = []
# Fixed size - multiple of four
arr = np.ones((20 + 2, 24 + 2)) * 5.0
f, g = ps.fields(f=arr, g=arr)
configurations.append((arr, f, g))
# Fixed size - no multiple of four
arr = np.ones((21 + 2, 25 + 2)) * 5.0
f, g = ps.fields(f=arr, g=arr)
configurations.append((arr, f, g))
# Fixed size - different remainder
arr = np.ones((23 + 2, 17 + 2)) * 5.0
f, g = ps.fields(f=arr, g=arr)
configurations.append((arr, f, g))
for arr, f, g in configurations:
update_rule = [ps.Assignment(g[0, 0], f[0, 0] + f[-1, 0] + f[1, 0] + f[0, 1] + f[0, -1] + 42.0)]
ast = ps.create_kernel(update_rule)
vectorize(ast)
func = ast.compile()
dst = np.zeros_like(arr)
func(g=dst, f=arr)
np.testing.assert_equal(dst[1:-1, 1:-1], 5 * 5.0 + 42.0)
def test_mean_filter_evaluation_gpu():
from pycuda.gpuarray import to_gpu
x, y = pystencils.fields('x,y: float32[2d]')
x_array = to_gpu(np.random.rand(20, 23, 25).astype(np.float32))
y_array = to_gpu(np.random.rand(20, 23, 25).astype(np.float32))
assignments = pystencils_reco.filters.mean_filter(x, y,
BoxStencil(3, ndim=2))
ast = pystencils.create_kernel(assignments, target='gpu')
kernel = ast.compile()
kernel(x=x_array[0], y=y_array[0])
kernel = pystencils_reco.filters.mean_filter(x, y, BoxStencil(
3, ndim=2)).compile(target='gpu')
kernel(x=x_array[0], y=y_array[0])
kernel = pystencils_reco.filters.mean_filter(
x, y, BoxStencil(3, ndim=2, with_center=False)).compile(target='gpu')
kernel(x=x_array[0], y=y_array[0])
x, y = pystencils.fields('x,y: float32[3d]')
kernel = pystencils_reco.filters.mean_filter(x, y, BoxStencil(
3, ndim=3)).compile(target='gpu')
kernel(x=x_array, y=y_array)
kernel = pystencils_reco.filters.mean_filter(
x, y, BoxStencil(3, ndim=3, with_center=False)).compile(target='gpu')
kernel(x=x_array, y=y_array)
def test_gauss_filter_evaluation():
x_array = np.random.rand(20, 23, 25).astype(np.float32)
y_array = np.random.rand(20, 23, 25).astype(np.float32)
x, y = pystencils.fields('x,y: float32[3d]')
kernel = pystencils_reco.filters.gauss_filter(x,
y,
BoxStencil(3, ndim=3),
sigma=0.5).compile()
kernel(x=x_array, y=y_array)
x, y = pystencils.fields('x,y: float32[2d]')
kernel = pystencils_reco.filters.gauss_filter(x,
y,
BoxStencil(
3,
ndim=2,
with_center=False),
sigma=0.5).compile()
kernel(x=x_array[0], y=y_array[0])
kernel = pystencils_reco.filters.gauss_filter(
x,
y,
BoxStencil(3, ndim=2, with_center=False),
sigma=sympy.Symbol('a')).compile()
kernel(x=x_array[0], y=y_array[0], a=0.7)
def test_packinfo_walberla_gen():
for openmp in (False, True):
for da in (False, True):
with ManualCodeGenerationContext(openmp=openmp,
double_accuracy=da) as ctx:
dtype = "float64" if ctx.double_accuracy else "float32"
f, g = ps.fields("f, g(4): {}[3D]".format(dtype))
generate_pack_info_for_field(ctx, 'PI1', f)
src, dst = ps.fields("src, src_tmp: {}[2D]".format(dtype))
stencil = [[0, -1, 0], [-1, 4, -1], [0, -1, 0]]
assignments = [
ps.assignment_from_stencil(stencil,
src,
dst,
normalization_factor=4)
]
generate_pack_info_from_kernel(ctx, 'PI2', assignments)
expected_files = ('PI1.cpp', 'PI1.h', 'PI2.cpp', 'PI2.h')
assert all(e in ctx.files for e in expected_files)
for file_name_to_test in ('PI1.cpp', 'PI2.cpp'):
file_to_test = ctx.files[file_name_to_test]
if openmp:
assert '#pragma omp parallel' in file_to_test
if da:
assert 'float ' not in file_to_test
else:
assert 'double ' not in file_to_test
def pystencils_2d_cpu_impl(x, y, coef, N, I=1):
if x.dtype == np.dtype('f4'):
src, dst = ps.fields(
'src, dst: float32[2D]',
src=x, dst=y
)
elif x.dtype == np.dtype('f8'):
src, dst = ps.fields(
'src, dst: double[2D]',
src=x, dst=y
)
else:
raise TypeError
if N == 1:
update_rule = make_2d_update_rule_1(src, dst, coef)
elif N == 6:
update_rule = make_2d_update_rule_6(src, dst, coef)
else:
raise ValueError
kernel = ps.create_kernel(update_rule, cpu_openmp=True).compile()
s = time.time()
for i in range(I):
if i % 2 == 0:
kernel(src=x, dst=y)
else:
kernel(src=y, dst=x)
e = time.time()
if (I - 1) % 2 == 0:
res = y
else:
res = x
return e - s, res
def test_sqrt_of_integer():
"""Regression test for bug where sqrt(3) was classified as integer"""
f = ps.fields("f: [1D]")
tmp = sp.symbols("tmp")
assignments = [ps.Assignment(tmp, sp.sqrt(3)),
ps.Assignment(f[0], tmp)]
arr_double = np.array([1], dtype=np.float64)
kernel = ps.create_kernel(assignments).compile()
kernel(f=arr_double)
assert 1.7 < arr_double[0] < 1.8
f = ps.fields("f: float32[1D]")
tmp = sp.symbols("tmp")
assignments = [ps.Assignment(tmp, sp.sqrt(3)),
ps.Assignment(f[0], tmp)]
arr_single = np.array([1], dtype=np.float32)
config = ps.CreateKernelConfig(data_type="float32")
kernel = ps.create_kernel(assignments, config=config).compile()
kernel(f=arr_single)
code = ps.get_code_str(kernel.ast)
# ps.show_code(kernel.ast)
# 1.7320508075688772935 --> it is actually correct to round to ...773. This was wrong before !282
assert "1.7320508075688773f" in code
assert 1.7 < arr_single[0] < 1.8
def test_projection():
volume = pystencils.fields('volume: float32[100,200,300]')
projections = pystencils.fields('projections: float32[2D]')
projection_matrix = sympy.Matrix(
[[-289.0098977737411, -1205.2274801832275, 0.0, 186000.0],
[-239.9634468375339, -4.188577544948043, 1200.0, 144000.0],
[-0.9998476951563913, -0.01745240643728351, 0.0, 600.0]])
assignments = forward_projection(
volume, projections, projection_matrix) # .compile(target='gpu')
print(type(assignments))
print(pickle.dumps(assignments))
# a = sympy.Symbol('a')
# projection_matrix = sympy.Matrix([[-289.0098977737411, -1205.2274801832275, 0.0, 186000.0],
# [a, - 4.188577544948043, 1200.0, 144000.0],
# [-0.9998476951563913, -0.01745240643728351, 0.0, 600.0]])
# kernel = forward_projection(volume, projections, projection_matrix).compile(target='gpu')
# print(kernel.code)
a = sympy.symbols('a:12')
a = [pystencils.TypedSymbol(s.name, 'float32') for s in a]
A = sympy.Matrix(3, 4, lambda i, j: a[i * 4 + j])
# A = sympy.MatrixSymbol('A', 3, 4)
projection_matrix = A
forward_projection(volume, projections,
projection_matrix).compile(target='gpu')
def test_loop_independence_checks():
f, g = fields("f, g : double[2D]")
v = fields("v(2) : double[2D]")
with pytest.raises(ValueError) as e:
create_kernel(
[Assignment(g[0, 1], f[0, 1]),
Assignment(g[0, 0], f[1, 0])])
assert 'Field g is written at two different locations' in str(e.value)
# This is allowed - because only one element of g is accessed
create_kernel(
[Assignment(g[0, 2], f[0, 1]),
Assignment(g[0, 2], 2 * g[0, 2])])
create_kernel([
Assignment(v[0, 2](1), f[0, 1]),
Assignment(v[0, 1](0), 4),
Assignment(v[0, 2](1), 2 * v[0, 2](1))
])
with pytest.raises(ValueError) as e:
create_kernel(
[Assignment(g[0, 1], 3),
Assignment(f[0, 1], 2 * g[0, 2])])
assert 'Field g is read at (0, 2) and written at (0, 1)' in str(e.value)
def test_mean_filter_evaluation():
x, y = pystencils.fields('x,y: float32[2d]')
x_array = np.random.rand(20, 23, 25).astype(np.float32)
y_array = np.random.rand(20, 23, 25).astype(np.float32)
assignments = pystencils_reco.filters.mean_filter(x, y,
BoxStencil(3, ndim=2))
ast = pystencils.create_kernel(assignments)
kernel = ast.compile()
kernel(x=x_array[0], y=y_array[0])
kernel = pystencils_reco.filters.mean_filter(x, y,
BoxStencil(3,
ndim=2)).compile()
kernel(x=x_array[0], y=y_array[0])
kernel = pystencils_reco.filters.mean_filter(
x, y, BoxStencil(3, ndim=2, with_center=False)).compile()
kernel(x=x_array[0], y=y_array[0])
x, y = pystencils.fields('x,y: float32[3d]')
kernel = pystencils_reco.filters.mean_filter(x, y,
BoxStencil(3,
ndim=3)).compile()
kernel(x=x_array, y=y_array)
kernel = pystencils_reco.filters.mean_filter(
x, y, BoxStencil(3, ndim=3, with_center=False)).compile()
kernel(x=x_array, y=y_array)
def test_hydro_lb():
stencil_hydro = LBStencil(Stencil.D3Q27)
density_liquid = 1.0
density_gas = 0.001
surface_tension = 0.0001
W = 5
# phase-field parameter
drho3 = (density_liquid - density_gas) / 3
# coefficient related to surface tension
beta = 12.0 * (surface_tension / W)
# coefficient related to surface tension
kappa = 1.5 * surface_tension * W
u = fields("vel_field(" + str(stencil_hydro.D) + "): [" + str(stencil_hydro.D) + "D]", layout='fzyx')
C = fields("phase_field: [" + str(stencil_hydro.D) + "D]", layout='fzyx')
g = fields("lb_velocity_field(" + str(stencil_hydro.Q) + "): [" + str(stencil_hydro.D) + "D]", layout='fzyx')
g_tmp = fields("lb_velocity_field_tmp(" + str(stencil_hydro.Q) + "): [" + str(stencil_hydro.D) + "D]", layout='fzyx')
# calculate the relaxation rate for the hydro lb as well as the body force
density = density_gas + C.center * (density_liquid - density_gas)
# force acting on all phases of the model
body_force = [0, 0, 0]
relaxation_time = 0.03 + 0.5
relaxation_rate = 1.0 / relaxation_time
lbm_config = LBMConfig(stencil=stencil_hydro, method=Method.MRT,
weighted=True, relaxation_rates=[relaxation_rate, 1, 1, 1, 1, 1])
method_hydro = create_lb_method(lbm_config=lbm_config)
# create the kernels for the initialization of the g and h field
g_updates = initializer_kernel_hydro_lb(g, u, method_hydro)
force_g = hydrodynamic_force(g, C, method_hydro,
relaxation_time, density_liquid, density_gas, kappa, beta, body_force)
force_model_g = MultiphaseForceModel(force=force_g, rho=density)
hydro_lb_update_rule_normal = get_collision_assignments_hydro(lb_method=method_hydro,
density=density,
velocity_input=u,
force_model=force_model_g,
sub_iterations=2,
symbolic_fields={"symbolic_field": g,
"symbolic_temporary_field": g_tmp},
kernel_type='collide_only')
hydro_lb_update_rule_push = get_collision_assignments_hydro(lb_method=method_hydro,
density=density,
velocity_input=u,
force_model=force_model_g,
sub_iterations=2,
symbolic_fields={"symbolic_field": g,
"symbolic_temporary_field": g_tmp},
kernel_type='collide_stream_push')
def test_fields_from_torch_tensor():
torch = pytest.importorskip('torch')
import torch
a, b = torch.zeros((20, 10)), torch.zeros((6, 7))
x, y = pystencils.fields(x=a, y=b)
print(x)
print(y)
c = torch.zeros((20, 10)).cuda()
z = pystencils.fields(z=c)
print(z)
def test_projection_cpu():
volume = pystencils.fields('volume: float32[3d]')
projections = pystencils.fields('projections: float32[2D]')
projection_matrix = sympy.Matrix(
[[-289.0098977737411, -1205.2274801832275, 0.0, 186000.0],
[-239.9634468375339, -4.188577544948043, 1200.0, 144000.0],
[-0.9998476951563913, -0.01745240643728351, 0.0, 600.0]])
forward_projection(volume, projections, projection_matrix).compile()
def test_genereric_projection():
volume = pystencils.fields('volume: float32[3d]')
projections = pystencils.fields('projections: float32[2D]')
projection_matrix = pystencils_reco.matrix_symbols(
'T', pystencils.data_types.create_type('float32'), 3, 4)
assignments = forward_projection(volume, projections, projection_matrix)
kernel = assignments.compile('gpu')
pystencils.show_code(kernel)
def test_jacobi3d_fixed_size():
print("Fixed Size: Large non divisible sizes")
arr = np.empty([10, 10, 9])
src, dst = ps.fields("src, dst: double[3D]", src=arr, dst=arr)
check_equivalence(jacobi(dst, src), arr)
print("Fixed Size: Smaller than block sizes")
arr = np.empty([4, 5, 6])
src, dst = ps.fields("src, dst: double[3D]", src=arr, dst=arr)
check_equivalence(jacobi(dst, src), arr)
print("Fixed Size: Multiples of block sizes")
arr = np.empty([8*4, 16*2, 4*3])
src, dst = ps.fields("src, dst: double[3D]", src=arr, dst=arr)
check_equivalence(jacobi(dst, src), arr)
def test_codegen():
for openmp in (False, True):
for da in (False, True):
with ManualCodeGenerationContext(openmp=openmp,
double_accuracy=da) as ctx:
h = sp.symbols("h")
dtype = "float64" if ctx.double_accuracy else "float32"
# ----- Jacobi 2D - created by specifying weights in nested list --------------------------
src, dst = ps.fields("src, src_tmp: {}[2D]".format(dtype))
stencil = [[0, -1, 0], [-1, 4, -1], [0, -1, 0]]
assignments = ps.assignment_from_stencil(
stencil, src, dst, normalization_factor=4 * h**2)
generate_sweep(ctx,
'JacobiKernel2D',
assignments,
field_swaps=[(src, dst)])
# ----- Jacobi 3D - created by using kernel_decorator with assignments in '@=' format -----
src, dst = ps.fields("src, src_tmp: {}[3D]".format(dtype))
@ps.kernel
def kernel_func():
dst[0, 0,
0] @= (src[1, 0, 0] + src[-1, 0, 0] +
src[0, 1, 0] + src[0, -1, 0] +
src[0, 0, 1] + src[0, 0, -1]) / (6 * h**2)
generate_sweep(ctx,
'JacobiKernel3D',
kernel_func,
field_swaps=[(src, dst)])
expected_files = ('JacobiKernel3D.cpp', 'JacobiKernel3D.h',
'JacobiKernel2D.cpp', 'JacobiKernel2D.h')
assert all(e in ctx.files for e in expected_files)
for file_name_to_test in ('JacobiKernel3D.cpp',
'JacobiKernel2D.cpp'):
file_to_test = ctx.files[file_name_to_test]
if openmp:
assert '#pragma omp parallel' in file_to_test
if da:
assert 'float ' not in file_to_test
else:
assert 'double ' not in file_to_test
def test_vec_any(instruction_set, dtype):
if instruction_set in ['sve', 'rvv']:
width = 4 # we don't know the actual value
else:
width = get_vector_instruction_set(dtype, instruction_set)['width']
data_arr = np.zeros((4 * width, 4 * width),
dtype=np.float64 if dtype == 'double' else np.float32)
data_arr[3:9, 1:3 * width - 1] = 1.0
data = ps.fields(f"data: {dtype}[2D]", data=data_arr)
c = [
ps.Assignment(sp.Symbol("t1"), vec_any(data.center() > 0.0)),
Conditional(vec_any(data.center() > 0.0),
Block([ps.Assignment(data.center(), 2.0)]))
]
ast = ps.create_kernel(
c,
target=ps.Target.CPU,
cpu_vectorize_info={'instruction_set': instruction_set})
kernel = ast.compile()
kernel(data=data_arr)
if instruction_set in ['sve', 'rvv']:
# we only know that the first value has changed
np.testing.assert_equal(data_arr[3:9, :3 * width - 1], 2.0)
else:
np.testing.assert_equal(data_arr[3:9, :3 * width], 2.0)
def test_prod_var_limit():
k = pystencils.TypedSymbol('k', create_type('int64'))
limit = pystencils.TypedSymbol('limit', create_type('int64'))
sum = sympy.Sum(k, (k, 1, limit))
expanded_sum = sum.replace(limit, 100).doit()
print(sum)
print(expanded_sum)
x = pystencils.fields('x: int64[1d]')
assignments = pystencils.AssignmentCollection({x.center(): sum})
ast = pystencils.create_kernel(assignments)
code = str(pystencils.show_code(ast))
kernel = ast.compile()
print(code)
array = np.zeros((10, ), np.int64)
kernel(x=array, limit=100)
assert np.allclose(array, int(expanded_sum) * np.ones_like(array))
def test_sum_use_float():
sum = sympy.Sum(k, (k, 1, 100))
expanded_sum = sum.doit()
print(sum)
print(expanded_sum)
x = pystencils.fields('x: float32[1d]')
assignments = pystencils.AssignmentCollection({x.center(): sum})
ast = pystencils.create_kernel(assignments,
data_type=create_type('float32'))
code = str(pystencils.show_code(ast))
kernel = ast.compile()
print(code)
print(pystencils.show_code(ast))
assert 'float sum' in code
array = np.zeros((10, ), np.float32)
kernel(x=array)
assert np.allclose(array, int(expanded_sum) * np.ones_like(array))
def test_tensorflow_jit_cpu():
pytest.importorskip('tensorflow')
module_name = "Ololol"
target = 'cpu'
z, y, x = pystencils.fields("z, y, x: [20,40]")
a = sympy.Symbol('a')
forward_assignments = pystencils.AssignmentCollection(
{z[0, 0]: x[0, 0] * sympy.log(a * x[0, 0] * y[0, 0])})
backward_assignments = create_backward_assignments(forward_assignments)
forward_ast = pystencils.create_kernel(forward_assignments, target)
forward_ast.function_name = 'forward_jit'
backward_ast = pystencils.create_kernel(backward_assignments, target)
backward_ast.function_name = 'backward_jit'
module = TensorflowModule(module_name, [forward_ast, backward_ast])
lib = pystencils_autodiff.tensorflow_jit.compile_sources_and_load(
[str(module)])
assert 'call_forward_jit' in dir(lib)
assert 'call_backward_jit' in dir(lib)
lib = module.compile()
assert 'call_forward_jit' in dir(lib)
assert 'call_backward_jit' in dir(lib)
def test_product(default_assignment_simplifications):
k = ps.TypedSymbol('k', create_type('int64'))
sum = sympy.Product(k, (k, 1, 10))
expanded_sum = sum.doit()
print(sum)
print(expanded_sum)
x = ps.fields('x: int64[1d]')
assignments = ps.AssignmentCollection({x.center(): sum})
config = ps.CreateKernelConfig(
default_assignment_simplifications=default_assignment_simplifications)
ast = ps.create_kernel(assignments, config=config)
code = ps.get_code_str(ast)
kernel = ast.compile()
print(code)
if default_assignment_simplifications is False:
assert 'int64_t product' in code
array = np.zeros((10, ), np.int64)
kernel(x=array)
assert np.allclose(array, int(expanded_sum) * np.ones_like(array))
def test_symbol_renaming():
"""When two loops have assignments to the same symbol with different rhs and both
are pulled before the loops, one of them has to be renamed
"""
f, g = ps.fields("f, g : double[2D]")
a, b, c = [TypedSymbol(n, np.float64) for n in ('a', 'b', 'c')]
loop1 = LoopOverCoordinate(
Block(
[SympyAssignment(c, a + b),
SympyAssignment(g[0, 0], f[0, 0] + c)]), 0, 0, 10)
loop2 = LoopOverCoordinate(
Block([
SympyAssignment(c, a**2 + b**2),
SympyAssignment(g[0, 0], f[0, 0] + c)
]), 0, 0, 10)
block = Block([loop1, loop2])
move_constants_before_loop(block)
loops = block.atoms(LoopOverCoordinate)
assert len(loops) == 2
for loop in loops:
assert len(loop.body.args) == 1
assert len(loop.parent.args) == 4 # 2 loops + 2 subexpressions
assert loop.parent.args[0].lhs.name != loop.parent.args[1].lhs.name
def test_Basic_data_type():
assert typed_symbols(("s", "f"), np.uint) == typed_symbols("s, f", np.uint)
t_symbols = typed_symbols(("s", "f"), np.uint)
s = t_symbols[0]
assert t_symbols[0] == TypedSymbol("s", np.uint)
assert s.dtype.is_uint()
assert s.dtype.is_complex() == 0
assert typed_symbols("s", str).dtype.is_other()
assert typed_symbols("s", bool).dtype.is_other()
assert typed_symbols("s", np.void).dtype.is_other()
assert typed_symbols("s", np.float64).dtype.base_name == 'double'
# removed for old sympy version
# assert typed_symbols(("s"), np.float64).dtype.sympy_dtype == typed_symbols(("s"), float).dtype.sympy_dtype
f, g = ps.fields("f, g : double[2D]")
expr = ps.Assignment(f.center(), 2 * g.center() + 5)
new_expr = type_all_numbers(expr, np.float64)
assert "cast_func(2, double)" in str(new_expr)
assert "cast_func(5, double)" in str(new_expr)
m = matrix_symbols("a, b", np.uint, 3, 3)
assert len(m) == 2
m = m[0]
for i, elem in enumerate(m):
assert elem == TypedSymbol(f"a{i}", np.uint)
assert elem.dtype.is_uint()
assert TypedSymbol("s", np.uint).canonical == TypedSymbol("s", np.uint)
assert TypedSymbol("s", np.uint).reversed == TypedSymbol("s", np.uint)
def test_vec_all(instruction_set, dtype):
if instruction_set in ['sve', 'rvv']:
width = 1000 # we don't know the actual value, need something guaranteed larger than vector
else:
width = get_vector_instruction_set(dtype, instruction_set)['width']
data_arr = np.zeros((4 * width, 4 * width),
dtype=np.float64 if dtype == 'double' else np.float32)
data_arr[3:9, 1:3 * width - 1] = 1.0
data = ps.fields(f"data: {dtype}[2D]", data=data_arr)
c = [
Conditional(vec_all(data.center() > 0.0),
Block([ps.Assignment(data.center(), 2.0)]))
]
ast = ps.create_kernel(
c,
target=Target.CPU,
cpu_vectorize_info={'instruction_set': instruction_set})
kernel = ast.compile()
kernel(data=data_arr)
if instruction_set in ['sve', 'rvv']:
# we only know that some values in the middle have been replaced
assert np.all(data_arr[3:9, :2] <= 1.0)
assert np.any(data_arr[3:9, 2:] == 2.0)
else:
np.testing.assert_equal(data_arr[3:9, :1], 0.0)
np.testing.assert_equal(data_arr[3:9, 1:width], 1.0)
np.testing.assert_equal(data_arr[3:9, width:2 * width], 2.0)
np.testing.assert_equal(data_arr[3:9, 2 * width:3 * width - 1], 1.0)
np.testing.assert_equal(data_arr[3:9, 3 * width - 1:], 0.0)
def test_simple_2d_check_assignment_collection():
# use simply example
z, y, x = ps.fields("z, y, x: [2d]")
forward_assignments = ps.AssignmentCollection([ps.Assignment(
z[0, 0], x[0, 0]*sp.log(x[0, 0]*y[0, 0]))], [])
jac = pystencils_autodiff.get_jacobian_of_assignments(
forward_assignments, [x[0, 0], y[0, 0]])
assert jac.shape == (len(forward_assignments.bound_symbols),
len(forward_assignments.free_symbols))
print(repr(jac))
assert repr(jac) == 'Matrix([[log(x_C*y_C) + 1, x_C/y_C]])'
for diff_mode in DiffModes:
pystencils_autodiff.create_backward_assignments(
forward_assignments, diff_mode=diff_mode)
pystencils_autodiff.create_backward_assignments(
pystencils_autodiff.create_backward_assignments(forward_assignments), diff_mode=diff_mode)
result1 = pystencils_autodiff.create_backward_assignments(
forward_assignments, diff_mode=DiffModes.TRANSPOSED)
result2 = pystencils_autodiff.create_backward_assignments(
forward_assignments, diff_mode=DiffModes.TF_MAD)
assert result1 == result2
def test_sum_use_float(default_assignment_simplifications):
sum = sympy.Sum(sp.abc.k, (sp.abc.k, 1, 100))
expanded_sum = sum.doit()
print(sum)
print(expanded_sum)
x = ps.fields('x: float32[1d]')
assignments = ps.AssignmentCollection({x.center(): sum})
config = ps.CreateKernelConfig(
default_assignment_simplifications=default_assignment_simplifications,
data_type=create_type('float32'))
ast = ps.create_kernel(assignments, config=config)
code = ps.get_code_str(ast)
kernel = ast.compile()
print(code)
if default_assignment_simplifications is False:
assert 'float sum' in code
array = np.zeros((10, ), np.float32)
kernel(x=array)
assert np.allclose(array, int(expanded_sum) * np.ones_like(array))
def test_assignment_collection_dict_conversion():
x, y = pystencils.fields('x,y: [2D]')
collection_normal = pystencils.AssignmentCollection(
[pystencils.Assignment(x.center(), y[1, 0] + y[0, 0])], [])
collection_dict = pystencils.AssignmentCollection(
{x.center(): y[1, 0] + y[0, 0]}, {})
assert str(collection_normal) == str(collection_dict)
assert collection_dict.main_assignments_dict == {
x.center(): y[1, 0] + y[0, 0]
}
assert collection_dict.subexpressions_dict == {}
collection_normal = pystencils.AssignmentCollection([
pystencils.Assignment(y[1, 0], x.center()),
pystencils.Assignment(y[0, 0], x.center())
], [])
collection_dict = pystencils.AssignmentCollection(
{
y[1, 0]: x.center(),
y[0, 0]: x.center()
}, {})
assert str(collection_normal) == str(collection_dict)
assert collection_dict.main_assignments_dict == {
y[1, 0]: x.center(),
y[0, 0]: x.center()
}
assert collection_dict.subexpressions_dict == {}
def test_module_printing_parameter():
module_name = "Ololol"
for target in ('cpu', 'gpu'):
z, y, x = pystencils.fields("z, y, x: [20,40]")
a = sympy.Symbol('a')
forward_assignments = pystencils.AssignmentCollection(
{z[0, 0]: x[0, 0] * sympy.log(a * x[0, 0] * y[0, 0])})
backward_assignments = create_backward_assignments(forward_assignments)
forward_ast = pystencils.create_kernel(forward_assignments, target)
forward_ast.function_name = 'forward'
backward_ast = pystencils.create_kernel(backward_assignments, target)
backward_ast.function_name = 'backward'
module = TorchModule(module_name, [forward_ast, backward_ast])
print(module)
module = TensorflowModule(module_name, {forward_ast: backward_ast})
print(module)
if target == 'cpu':
module = PybindModule(module_name, [forward_ast, backward_ast])
print(module)
module = PybindModule(module_name, forward_ast)
print(module)
def test_tfmad_gradient_check_torch():
torch = pytest.importorskip('torch')
a, b, out = ps.fields("a, b, out: float[5,7]")
cont = 2 * ps.fd.Diff(a, 0) - 1.5 * ps.fd.Diff(a, 1) \
- ps.fd.Diff(b, 0) + 3 * ps.fd.Diff(b, 1)
discretize = ps.fd.Discretization2ndOrder(dx=1)
discretization = discretize(cont) + 1.2 * a.center
assignment = ps.Assignment(out.center(), discretization)
assignment_collection = ps.AssignmentCollection([assignment], [])
print('Forward')
print(assignment_collection)
print('Backward')
auto_diff = pystencils_autodiff.AutoDiffOp(assignment_collection,
diff_mode='transposed-forward')
backward = auto_diff.backward_assignments
print(backward)
print('Forward output fields (to check order)')
print(auto_diff.forward_input_fields)
a_tensor = torch.zeros(*a.shape, dtype=torch.float64, requires_grad=True)
b_tensor = torch.zeros(*b.shape, dtype=torch.float64, requires_grad=True)
function = auto_diff.create_tensorflow_op({
a: a_tensor,
b: b_tensor
},
backend='torch')
torch.autograd.gradcheck(function.apply, [a_tensor, b_tensor])
def test_rotation_compilation():
for ndim in (2, 3):
for angle in (0.4, -0.8):
for axis in range(3):
x, y = pystencils.fields('x,y: float32[%id]' % ndim)
rotation_transform(x, y, angle, axis).compile('gpu')
def test_scaling():
for ndim in range(1, 5):
for scale in (0.5, [(s + 1) * 0.1 for s in range(ndim)]):
x, y = pystencils.fields('x,y: float32[%id]' % ndim)
transform = scale_transform(x, y, scale)
print(transform)
import sympy as sp
import pystencils as ps
from pystencils_walberla import CodeGeneration, generate_sweep
with CodeGeneration() as ctx:
h = sp.symbols("h")
# ----- Jacobi 2D - created by specifying weights in nested list --------------------------
src, dst = ps.fields("src, src_tmp: [2D]", layout='fzyx')
stencil = [[0, 1, 0],
[1, 0, 1],
[0, 1, 0]]
assignments = ps.assignment_from_stencil(stencil, src, dst, normalization_factor=1 / (4 * h ** 2))
generate_sweep(ctx, 'JacobiKernel2D', assignments, field_swaps=[(src, dst)])
# ----- Jacobi 3D - created by using kernel_decorator with assignments in '@=' format -----
src, dst = ps.fields("src, src_tmp: [3D]")
@ps.kernel
def kernel_func():
dst[0, 0, 0] @= (src[1, 0, 0] + src[-1, 0, 0] +
src[0, 1, 0] + src[0, -1, 0] +
src[0, 0, 1] + src[0, 0, -1]) / (h ** 2 * 6)
generate_sweep(ctx, 'JacobiKernel3D', kernel_func, field_swaps=[(src, dst)])
import sympy as sp
import pystencils as ps
from pystencils_walberla import CodeGeneration, generate_sweep
with CodeGeneration() as ctx:
h = sp.symbols("h")
# ----- Jacobi 2D - created by specifying weights in nested list --------------------------
src, dst = ps.fields("src, src_tmp: [2D]")
stencil = [[0, -1, 0],
[-1, 4, -1],
[0, -1, 0]]
assignments = ps.assignment_from_stencil(stencil, src, dst, normalization_factor=4 * h**2)
generate_sweep(ctx, 'CudaJacobiKernel2D', assignments, field_swaps=[(src, dst)], target="gpu")
# ----- Jacobi 3D - created by using kernel_decorator with assignments in '@=' format -----
src, dst = ps.fields("src, src_tmp: [3D]")
@ps.kernel
def kernel_func():
dst[0, 0, 0] @= (src[1, 0, 0] + src[-1, 0, 0] +
src[0, 1, 0] + src[0, -1, 0] +
src[0, 0, 1] + src[0, 0, -1]) / (6 * h ** 2)
generate_sweep(ctx, 'CudaJacobiKernel3D', kernel_func, field_swaps=[(src, dst)], target="gpu")
import sympy as sp
import pystencils as ps
from lbmpy.creationfunctions import create_lb_method
from lbmpy.boundaries import NoSlip, UBB
from pystencils_walberla import CodeGeneration
from lbmpy_walberla import generate_lattice_model, RefinementScaling, generate_boundary
with CodeGeneration() as ctx:
omega = sp.Symbol("omega")
force_field = ps.fields("force(3): [3D]", layout='fzyx')
# lattice Boltzmann method
lb_method = create_lb_method(stencil='D3Q19', method='srt', relaxation_rates=[omega],
force_model='guo', force=force_field.center_vector)
scaling = RefinementScaling()
scaling.add_standard_relaxation_rate_scaling(omega)
scaling.add_force_scaling(force_field)
# generate components
generate_lattice_model(ctx, 'SrtWithForceFieldModel', lb_method, refinement_scaling=scaling)
generate_boundary(ctx, 'MyUBB', UBB([0.05, 0, 0]), lb_method)
generate_boundary(ctx, 'MyNoSlip', NoSlip(), lb_method)
import pystencils as ps
from lbmpy.updatekernels import create_stream_pull_only_kernel
from lbmpy.stencils import get_stencil
from pystencils_walberla import CodeGeneration, generate_sweep
with CodeGeneration() as ctx:
f_size = 19
dtype = 'float64' if ctx.double_accuracy else 'float32'
# Copy sweep
src, dst = ps.fields("src({f_size}), dst({f_size}) : {dtype}[3D]".format(dtype=dtype, f_size=f_size),
layout='fzyx')
copy_only = [ps.Assignment(dst(i), src(i)) for i in range(f_size)]
generate_sweep(ctx, 'MicroBenchmarkCopyKernel', copy_only,
target='gpu', gpu_indexing_params={'block_size': (128, 1, 1)})
# Stream-only sweep
stencil = get_stencil("D3Q19")
stream_only = create_stream_pull_only_kernel(stencil, src_field_name='src', dst_field_name='dst',
generic_field_type=dtype, generic_layout='fzyx')
generate_sweep(ctx, 'MicroBenchmarkStreamKernel', stream_only,
target='gpu', gpu_indexing_params={'block_size': (128, 1, 1)})
