diff --git a/pystencils/backends/cbackend.py b/pystencils/backends/cbackend.py
index aea77de2230d37a731c17a493a29844cfa1000f8..06eda124eb9c64b93fe63ab604cf19f3a7766bcf 100644
--- a/pystencils/backends/cbackend.py
+++ b/pystencils/backends/cbackend.py
@@ -9,6 +9,7 @@ from sympy.core import S
from sympy.core.cache import cacheit
from sympy.logic.boolalg import BooleanFalse, BooleanTrue
from sympy.functions.elementary.trigonometric import TrigonometricFunction, InverseTrigonometricFunction
+from sympy.functions.elementary.hyperbolic import HyperbolicFunction
from pystencils.astnodes import KernelFunction, LoopOverCoordinate, Node
from pystencils.cpu.vectorization import vec_all, vec_any, CachelineSize
@@ -494,10 +495,18 @@ class CustomSympyPrinter(CCodePrinter):
arg, data_type = expr.args
if isinstance(arg, sp.Number) and arg.is_finite:
return self._typed_number(arg, data_type)
- elif isinstance(arg, (sp.Pow, InverseTrigonometricFunction, TrigonometricFunction)) and data_type == BasicType('float32'):
- known = self.known_functions[arg.__class__.__name__]
+ elif isinstance(arg, (InverseTrigonometricFunction, TrigonometricFunction, HyperbolicFunction)) \
+ and data_type == BasicType('float32'):
+ known = self.known_functions[arg.__class__.__name__.lower()]
code = self._print(arg)
return code.replace(known, f"{known}f")
+ elif isinstance(arg, sp.Pow) and data_type == BasicType('float32'):
+ known = ['sqrt', 'cbrt', 'pow']
+ code = self._print(arg)
+ for k in known:
+ if k in code:
+ return code.replace(k, f'{k}f')
+ raise ValueError(f"{code} doesn't give {known=} function back.")
else:
return f"(({data_type})({self._print(arg)}))"
elif isinstance(expr, fast_division):
diff --git a/pystencils/typing/leaf_typing.py b/pystencils/typing/leaf_typing.py
index 224a824472fdd4358bb2c63391503f35c2b9fa75..20f92eabdf8afe88e039d165539d00335dcd95af 100644
--- a/pystencils/typing/leaf_typing.py
+++ b/pystencils/typing/leaf_typing.py
@@ -8,8 +8,8 @@ import sympy as sp
from sympy import Piecewise
from sympy.core.relational import Relational
from sympy.functions.elementary.piecewise import ExprCondPair
-from sympy.functions.elementary.trigonometric import TrigonometricFunction
-from sympy.functions.elementary.trigonometric import InverseTrigonometricFunction
+from sympy.functions.elementary.trigonometric import TrigonometricFunction, InverseTrigonometricFunction
+from sympy.functions.elementary.hyperbolic import HyperbolicFunction
from sympy.codegen import Assignment
from sympy.logic.boolalg import BooleanFunction
from sympy.logic.boolalg import BooleanAtom
@@ -206,7 +206,7 @@ class TypeAdder:
else:
new_args.append(a)
return expr.func(*new_args) if new_args else expr, collated_type
- elif isinstance(expr, (sp.Pow, InverseTrigonometricFunction, TrigonometricFunction)):
+ elif isinstance(expr, (sp.Pow, InverseTrigonometricFunction, TrigonometricFunction, HyperbolicFunction)):
args_types = [self.figure_out_type(arg) for arg in expr.args]
collated_type = collate_types([t for _, t in args_types])
new_args = [a if t.dtype_eq(collated_type) else CastFunc(a, collated_type) for a, t in args_types]
diff --git a/pystencils_tests/test_math_functions.py b/pystencils_tests/test_math_functions.py
new file mode 100644
index 0000000000000000000000000000000000000000..031e2320bc1ca0063168bc4c20ccf5642df875e7
--- /dev/null
+++ b/pystencils_tests/test_math_functions.py
@@ -0,0 +1,67 @@
+import pytest
+import sympy as sp
+import numpy as np
+import pystencils as ps
+
+
+@pytest.mark.parametrize('dtype', ["float64", "float32"])
+@pytest.mark.parametrize('func', [sp.Pow, sp.atan2])
+@pytest.mark.parametrize('target', [ps.Target.CPU, ps.Target.GPU])
+def test_two_arguments(dtype, func, target):
+ if target == ps.Target.GPU:
+ pytest.importorskip("pycuda")
+ dh = ps.create_data_handling(domain_size=(10, 10), periodicity=True, default_target=target)
+
+ x = dh.add_array('x', values_per_cell=1, dtype=dtype)
+ dh.fill("x", 0.0, ghost_layers=True)
+ y = dh.add_array('y', values_per_cell=1, dtype=dtype)
+ dh.fill("y", 1.0, ghost_layers=True)
+ z = dh.add_array('z', values_per_cell=1, dtype=dtype)
+ dh.fill("z", 2.0, ghost_layers=True)
+
+ config = ps.CreateKernelConfig(target=target)
+
+ # test sp.Max with one argument
+ up = ps.Assignment(x.center, func(y.center, z.center))
+ ast = ps.create_kernel(up, config=config)
+ code = ps.get_code_str(ast)
+ if dtype == 'float32':
+ assert func.__name__.lower() in code
+ kernel = ast.compile()
+
+ dh.all_to_gpu()
+ dh.run_kernel(kernel)
+ dh.all_to_cpu()
+
+ np.testing.assert_allclose(dh.gather_array("x")[0, 0], float(func(1.0, 2.0).evalf()))
+
+
+@pytest.mark.parametrize('dtype', ["float64", "float32"])
+@pytest.mark.parametrize('func', [sp.sin, sp.cos, sp.sinh, sp.cosh, sp.atan])
+@pytest.mark.parametrize('target', [ps.Target.CPU, ps.Target.GPU])
+def test_single_arguments(dtype, func, target):
+ if target == ps.Target.GPU:
+ pytest.importorskip("pycuda")
+ dh = ps.create_data_handling(domain_size=(10, 10), periodicity=True, default_target=target)
+
+ x = dh.add_array('x', values_per_cell=1, dtype=dtype)
+ dh.fill("x", 0.0, ghost_layers=True)
+ y = dh.add_array('y', values_per_cell=1, dtype=dtype)
+ dh.fill("y", 1.0, ghost_layers=True)
+
+ config = ps.CreateKernelConfig(target=target)
+
+ # test sp.Max with one argument
+ up = ps.Assignment(x.center, func(y.center))
+ ast = ps.create_kernel(up, config=config)
+ code = ps.get_code_str(ast)
+ if dtype == 'float32':
+ assert func.__name__.lower() in code
+ kernel = ast.compile()
+
+ dh.all_to_gpu()
+ dh.run_kernel(kernel)
+ dh.all_to_cpu()
+
+ np.testing.assert_allclose(dh.gather_array("x")[0, 0], float(func(1.0).evalf()),
+ rtol=10**-3 if dtype == 'float32' else 10**-5)
diff --git a/pystencils_tests/test_trigonometric_functions.py b/pystencils_tests/test_trigonometric_functions.py
deleted file mode 100644
index 24c33c7446ddd0b3a8ce0a1d1fef0b9b5c7964d9..0000000000000000000000000000000000000000
--- a/pystencils_tests/test_trigonometric_functions.py
+++ /dev/null
@@ -1,28 +0,0 @@
-import pytest
-import sympy as sp
-import numpy as np
-import math
-import pystencils as ps
-
-
-@pytest.mark.parametrize('dtype', ["float64", "float32"])
-def test_trigonometric_functions(dtype):
- dh = ps.create_data_handling(domain_size=(10, 10), periodicity=True)
-
- x = dh.add_array('x', values_per_cell=1, dtype=dtype)
- dh.fill("x", 0.0, ghost_layers=True)
- y = dh.add_array('y', values_per_cell=1, dtype=dtype)
- dh.fill("y", 1.0, ghost_layers=True)
- z = dh.add_array('z', values_per_cell=1, dtype=dtype)
- dh.fill("z", 2.0, ghost_layers=True)
-
- # config = pystencils.CreateKernelConfig(default_number_float=dtype)
-
- # test sp.Max with one argument
- up = ps.Assignment(x.center, sp.atan2(y.center, z.center))
- ast = ps.create_kernel(up)
- code = ps.get_code_str(ast)
- kernel = ast.compile()
- dh.run_kernel(kernel)
-
- np.testing.assert_allclose(dh.gather_array("x")[0, 0], math.atan2(1.0, 2.0))