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)})