diff --git a/src/copapy/__init__.py b/src/copapy/__init__.py index 941ea4c..dfe53fe 100644 --- a/src/copapy/__init__.py +++ b/src/copapy/__init__.py @@ -1,7 +1,7 @@ from ._target import Target from ._basic_types import NumLike, variable, generic_sdb, iif from ._vectors import vector -from ._math import sqrt, abs, sin, cos, tan, asin, acos, atan, atan2, log, exp, pow +from ._math import sqrt, abs, sin, cos, tan, asin, acos, atan, atan2, log, exp, pow, get_42 __all__ = [ "Target", @@ -21,5 +21,6 @@ __all__ = [ "atan2", "log", "exp", - "pow" + "pow", + "get_42" ] diff --git a/src/copapy/_math.py b/src/copapy/_math.py index dfa251c..02ac4ae 100644 --- a/src/copapy/_math.py +++ b/src/copapy/_math.py @@ -215,9 +215,15 @@ def acos(x: NumLike) -> variable[float] | float: return math.pi / 2 - asin(x) -def get_42() -> variable[float]: +@overload +def get_42(x: float | int) -> float: ... +@overload +def get_42(x: variable[Any]) -> variable[float]: ... +def get_42(x: NumLike) -> variable[float] | float: """Returns the variable representing the constant 42""" - return add_op('get_42', [0.0, 0.0]) + if isinstance(x, variable): + return add_op('get_42', [x, x]) + return float((int(x) * 3.0 + 42.0) * 5.0 + 21.0) def abs(x: T) -> T: diff --git a/stencils/aux_functions.c b/stencils/aux_functions.c index f82a668..a7ab664 100644 --- a/stencils/aux_functions.c +++ b/stencils/aux_functions.c @@ -11,6 +11,10 @@ int floor_div(float arg1, float arg2) { return i; } -float aux_get_42(float n) { - return n + 42.0; +NOINLINE float auxsub_get_42(int n) { + return n * 5.0f + 21.0f; } + +NOINLINE float aux_get_42(float n) { + return auxsub_get_42(n * 3.0f + 42.0f); +} \ No newline at end of file