test_reduction.py

import pytest
import numpy as np
import sympy as sp

import pystencils as ps
from pystencils import AddReducedAssignment


@pytest.mark.parametrize('dtype', ["float64", "float32"])
def test_log(dtype):
    a = sp.Symbol("a")
    x = ps.fields(f'x: {dtype}[1d]')

    # kernel with main assignments and no reduction

    main_assignment = ps.AssignmentCollection({x.center(): a})

    ast_main = ps.create_kernel(main_assignment, default_dtype=dtype)
    code_main = ps.get_code_str(ast_main)
    kernel_main = ast_main.compile()

    # ps.show_code(ast)

    if dtype == "float64":
        assert "float" not in code_main

    array = np.zeros((10,), dtype=dtype)
    kernel_main(x=array, a=100)
    assert np.allclose(array, 4.60517019)

    # kernel with single reduction assignment

    omega = sp.Symbol("omega")

    reduction_assignment = AddReducedAssignment(omega, x.center())

    ast_reduction = ps.create_kernel(reduction_assignment, default_dtype=dtype)
    code_reduction = ps.get_code_str(ast_reduction)
    kernel_reduction = ast_reduction.compile()

    if dtype == "float64":
        assert "float" not in code_reduction

    ps.show_code(ast_reduction)