diff --git a/pystencils/stencil.py b/pystencils/stencil.py index 359a0a29492c537957b8e24a952fa7e052ee898f..28d179c50d2716e4aa1530b2409abb77e6b10e66 100644 --- a/pystencils/stencil.py +++ b/pystencils/stencil.py @@ -34,6 +34,8 @@ def is_valid(stencil, max_neighborhood=None): True >>> is_valid([(2, 0), (1, 0)], max_neighborhood=1) False + >>> is_valid([(2, 0), (1, 0)], max_neighborhood=2) + True """ expected_dim = len(stencil[0]) for d in stencil: @@ -67,8 +69,11 @@ def have_same_entries(s1, s2): Examples: >>> stencil1 = [(1, 0), (-1, 0), (0, 1), (0, -1)] >>> stencil2 = [(-1, 0), (0, -1), (1, 0), (0, 1)] + >>> stencil3 = [(-1, 0), (0, -1), (1, 0)] >>> have_same_entries(stencil1, stencil2) True + >>> have_same_entries(stencil1, stencil3) + False """ if len(s1) != len(s2): return False diff --git a/pystencils_tests/test_staggered_kernel.py b/pystencils_tests/test_staggered_kernel.py index c4f491393d2b4d08fbb9b4c99356abdc7e7199b6..eecbf706c5538ad544d29a6b37366944aa504680 100644 --- a/pystencils_tests/test_staggered_kernel.py +++ b/pystencils_tests/test_staggered_kernel.py @@ -4,6 +4,7 @@ import sympy as sp import pytest import pystencils as ps +from pystencils import x_staggered_vector, TypedSymbol class TestStaggeredDiffusion: @@ -96,3 +97,12 @@ def test_staggered_loop_cutting(): assignments = [ps.Assignment(j.staggered_access("SW"), 1)] ast = ps.create_staggered_kernel(assignments, target=dh.default_target) assert not ast.atoms(ps.astnodes.Conditional) + + +def test_staggered_vector(): + dim = 2 + v = x_staggered_vector(dim) + ctr0 = TypedSymbol('ctr_0', 'int', nonnegative=True) + ctr1 = TypedSymbol('ctr_1', 'int', nonnegative=True) + expected_result = sp.Matrix(tuple((ctr0 + 0.5, ctr1 + 0.5))) + assert v == expected_result \ No newline at end of file diff --git a/pystencils_tests/test_stencils.py b/pystencils_tests/test_stencils.py index bf1418e053d41ebeb80b6ae623f1834b24a8a713..fa205ad7f04a446e1e16d5ea95e679cea70ecde4 100644 --- a/pystencils_tests/test_stencils.py +++ b/pystencils_tests/test_stencils.py @@ -1 +1,30 @@ import pystencils as ps +import sympy as sp + +from pystencils.stencil import coefficient_list, plot_expression + + +def test_coefficient_list(): + f = ps.fields("f: double[1D]") + expr = 2 * f[1] + 3 * f[-1] + coff = coefficient_list(expr) + assert coff == [3, 0, 2] + plot_expression(expr, matrix_form=True) + + f = ps.fields("f: double[3D]") + expr = 2 * f[1, 0, 0] + 3 * f[0, -1, 0] + coff = coefficient_list(expr) + assert coff == [[[0, 3, 0], [0, 0, 2], [0, 0, 0]]] + + expr = 2 * f[1, 0, 0] + 3 * f[0, -1, 0] + 4 * f[0, 0, 1] + coff = coefficient_list(expr, matrix_form=True) + assert coff[0] == sp.zeros(3, 3) + + # in 3D plot only works if there are entries on every of the three 2D planes. In the above examples z-1 was empty + expr = 2 * f[1, 0, 0] + 1 * f[0, -1, 0] + 1 * f[0, 0, 1] + f[0, 0, -1] + plot_expression(expr) + + +def test_plot_expression(): + f = ps.fields("f: double[2D]") + plot_expression(2 * f[1, 0] + 3 * f[0, -1], matrix_form=True) \ No newline at end of file