Merge pull request #172 from finsberg/generate-consistent-floats

Generate consistent floats

Author: finsberg

src/gotranx/cli/__init__.py

@@ -470,6 +470,11 @@ def ode2julia(
         1e-8,
         help="Delta value for the rush larsen schemes",
     ),
+    type_stable: bool = typer.Option(
+        False,
+        "--type-stable",
+        help="Add T to the function signature",
+    ),
     # format: CFormat = typer.Option(
     #     CFormat.clang_format,
     #     "--format",
@@ -497,6 +502,7 @@ def ode2julia(
         verbose=verbose,
         stiff_states=stiff_states,
         delta=delta,
+        type_stable=type_stable,
     )
 
 

src/gotranx/cli/gotran2julia.py

@@ -21,6 +21,7 @@ def get_code(
     missing_values: dict[str, int] | None = None,
     delta: float = 1e-8,
     stiff_states: list[str] | None = None,
+    type_stable: bool = False,
 ) -> str:
     """Generate the Julia code for the ODE
 
@@ -41,13 +42,17 @@ def get_code(
     stiff_states : list[str] | None, optional
         Stiff states, by default None. Only applicable for
         the hybrid rush larsen scheme
+    type_stable : bool, optional
+        Add T to the function signature, by default False
 
     Returns
     -------
     str
         The Julia code
     """
-    codegen = JuliaCodeGenerator(ode, remove_unused=remove_unused)  # , format=Format.none)
+    codegen = JuliaCodeGenerator(
+        ode, remove_unused=remove_unused, type_stable=type_stable
+    )  # , format=Format.none)
     # formatter = get_formatter(format=format)
 
     if missing_values is not None:
@@ -95,6 +100,7 @@ def main(
     missing_values: dict[str, int] | None = None,
     delta: float = 1e-8,
     stiff_states: list[str] | None = None,
+    type_stable: bool = False,
 ) -> None:
     loglevel = logging.DEBUG if verbose else logging.INFO
     structlog.configure(
@@ -109,6 +115,7 @@ def main(
         missing_values=missing_values,
         delta=delta,
         stiff_states=stiff_states,
+        type_stable=type_stable,
     )
     out = fname if outname is None else Path(outname)
     out_name = out.with_suffix(suffix=".jl")

src/gotranx/codegen/base.py

@@ -25,6 +25,7 @@ class Func(typing.NamedTuple):
     values_type: str
     return_name: str = "values"
     num_return_values: int = 0
+    post_function_signature: str = ""
 
 
 class RHSArgument(str, Enum):
@@ -332,6 +333,7 @@ def rhs(self, order: RHSArgument | str = RHSArgument.tsp, use_cse=False) -> str:
             shape_info="",
             values_type=rhs.values_type,
             missing_variables=missing_variables,
+            post_function_signature=rhs.post_function_signature,
         )
 
         return self._format(code)
@@ -400,6 +402,7 @@ def monitor_values(self, order: RHSArgument | str = RHSArgument.tsp, use_cse=Fal
             shape_info=shape_info,
             values_type="numpy.zeros(shape)",
             missing_variables=missing_variables,
+            post_function_signature=rhs.post_function_signature,
         )
 
         return self._format(code)
@@ -450,6 +453,7 @@ def missing_values(
             shape_info=shape_info,
             values_type="numpy.zeros(shape)",
             missing_variables=missing_variables,
+            post_function_signature=rhs.post_function_signature,
         )
 
         return self._format(code)
@@ -503,6 +507,7 @@ def scheme(self, f: schemes.scheme_func, order=SchemeArgument.stdp, **kwargs) ->
             shape_info="",
             values_type=rhs.values_type,
             missing_variables=missing_variables,
+            post_function_signature=rhs.post_function_signature,
         )
         return self._format(code)
 

src/gotranx/codegen/julia.py

@@ -16,12 +16,16 @@ def bool_to_int(expr: str) -> str:
 
 
 class GotranJuliaCodePrinter(JuliaCodePrinter):
-    def __init__(self, *args, **kwargs):
+    def __init__(self, type_stable: bool = False, *args, **kwargs):
         super().__init__(*args, **kwargs)
+        self._type_stable = type_stable
         self._settings["contract"] = False
 
     def _print_Float(self, flt):
-        return self._print(str(float(flt)))
+        value = str(float(flt))
+        if self._type_stable:
+            return self._print(f"T({value})")
+        return self._print(value)
 
     def _print_Piecewise(self, expr):
         if isinstance(expr.args[0][0], Assignment):
@@ -56,9 +60,9 @@ def _print_Indexed(self, expr):
 
 
 class JuliaCodeGenerator(CodeGenerator):
-    def __init__(self, ode: ODE, remove_unused: bool = False) -> None:
+    def __init__(self, ode: ODE, remove_unused: bool = False, type_stable: bool = False) -> None:
         super().__init__(ode, remove_unused=remove_unused)
-        self._printer = GotranJuliaCodePrinter()
+        self._printer = GotranJuliaCodePrinter(type_stable=type_stable)
         # setattr(self, "_formatter", get_formatter(format=format))
 
     @property
@@ -84,12 +88,23 @@ def _rhs_arguments(
         self, order: RHSArgument | str = RHSArgument.stp, const_states: bool = True
     ) -> Func:
         value = RHSArgument.get_value(order)
-        argument_dict = {
-            "s": "states",
-            "t": "t",
-            "p": "parameters",
-        }
-        argument_list = [argument_dict[v] for v in value] + ["values"]
+        if self._printer._type_stable:
+            argument_dict = {
+                "s": "states::AbstractVector{T}",
+                "t": "t::T",
+                "p": "parameters::AbstractVector{T}",
+            }
+            values = ["values::AbstractVector{T}"]
+            post_function_signature = " where T"
+        else:
+            argument_dict = {
+                "s": "states",
+                "t": "t",
+                "p": "parameters",
+            }
+            values = ["values"]
+            post_function_signature = ""
+        argument_list = [argument_dict[v] for v in value] + values
         states = sympy.IndexedBase("states", shape=(self.ode.num_states,), offset=1)
         parameters = sympy.IndexedBase("parameters", shape=(self.ode.num_parameters,), offset=1)
         values = sympy.IndexedBase("values", shape=(self.ode.num_states,), offset=1)
@@ -100,6 +115,7 @@ def _rhs_arguments(
             parameters=parameters,
             values=values,
             values_type="",
+            post_function_signature=post_function_signature,
         )
 
     def _scheme_arguments(
@@ -108,13 +124,26 @@ def _scheme_arguments(
         const_states: bool = True,
     ) -> Func:
         value = SchemeArgument.get_value(order)
-        argument_dict = {
-            "s": "states",
-            "t": "t",
-            "d": "dt",
-            "p": "parameters",
-        }
-        argument_list = [argument_dict[v] for v in value] + ["values"]
+        if self._printer._type_stable:
+            argument_dict = {
+                "s": "states::AbstractVector{T}",
+                "t": "t::T",
+                "d": "dt::T",
+                "p": "parameters::AbstractVector{T}",
+            }
+            values = ["values::AbstractVector{T}"]
+            post_function_signature = " where {T}"
+        else:
+            argument_dict = {
+                "s": "states",
+                "t": "t",
+                "d": "dt",
+                "p": "parameters",
+            }
+            values = ["values"]
+            post_function_signature = ""
+
+        argument_list = [argument_dict[v] for v in value] + values
         states = sympy.IndexedBase("states", shape=(self.ode.num_states,), offset=1)
         parameters = sympy.IndexedBase("parameters", shape=(self.ode.num_parameters,), offset=1)
         values = sympy.IndexedBase("values", shape=(self.ode.num_states,), offset=1)
@@ -125,4 +154,5 @@ def _scheme_arguments(
             parameters=parameters,
             values=values,
             values_type="",
+            post_function_signature=post_function_signature,
         )

src/gotranx/templates/__init__.py

@@ -151,6 +151,7 @@ def method(
         shape_info: str,
         values_type: str,
         missing_variables: str,
+        post_function_signature: str,
     ) -> str:
         """The method function is a function that generates a method
         for the model.
@@ -183,6 +184,8 @@ def method(
             The type of the values
         missing_variables : str
             The code for handling missing variables
+        post_function_signature : str
+            The code going after the function signature (e.g. 'where T' or '-> None')
 
         Returns
         -------

src/gotranx/templates/julia.py

@@ -45,14 +45,16 @@ def method(
     values: str,
     return_name: None = None,
     num_return_values: int = 0,
+    post_function_signature: str = "",
     **kwargs,
 ):
     indent_states = indent(states, "    ")
     indent_parameters = indent(parameters, "    ")
     indent_values = indent(values, "    ")
+
     return dedent(
         f"""
-function {name}({args})
+function {name}({args}){post_function_signature}
 
     # Assign states
 {indent_states}

tests/test_julia_codegen.py

@@ -191,3 +191,60 @@ def test_julia_monitored(codegen: JuliaCodeGenerator):
         "\nend"
         "\n"
     )
+
+
+def test_consistent_floats(parser, trans):
+    expr = """
+    \nstates(x=0)
+    \ndx_dt = Conditional(Ge(x, 31.4978), 1.0, 1.0763*exp(-1.007*exp(-0.0829*x)))
+    """
+
+    tree = parser.parse(expr)
+    result = trans.transform(tree)
+    ode = make_ode(*result, name="name")
+    codegen = JuliaCodeGenerator(ode)
+    rhs = codegen.rhs()
+
+    assert rhs == (
+        "\nfunction rhs(t, states, parameters, values)"
+        "\n"
+        "\n    # Assign states"
+        "\n    x = states[1]"
+        "\n"
+        "\n    # Assign parameters"
+        "\n"
+        "\n"
+        "\n    # Assign expressions"
+        "\n    dx_dt = ((x >= 31.4978) ? (1.0) : (1.0763 * exp((-1.007) * exp((-0.0829) * x))))"
+        "\n    values[1] = dx_dt"
+        "\nend"
+        "\n"
+    )
+
+
+def test_consistent_floats_with_T(parser, trans):
+    expr = """
+    \nstates(x=0)
+    \ndx_dt = Conditional(Ge(x, 31.4978), 1.0, 1.0763*exp(-1.007*exp(-0.0829*x)))
+    """
+
+    tree = parser.parse(expr)
+    result = trans.transform(tree)
+    ode = make_ode(*result, name="name")
+    codegen = JuliaCodeGenerator(ode, type_stable=True)
+    rhs = codegen.rhs()
+    assert rhs == (
+        "\nfunction rhs(t::T, states::AbstractVector{T}, parameters::AbstractVector{T}, values::AbstractVector{T}) where T"  # noqa: E501
+        "\n"
+        "\n    # Assign states"
+        "\n    x = states[1]"
+        "\n"
+        "\n    # Assign parameters"
+        "\n"
+        "\n"
+        "\n    # Assign expressions"
+        "\n    dx_dt = ((x >= T(31.4978)) ? (T(1.0)) : (T(1.0763) * exp((-T(1.007)) * exp((-T(0.0829)) * x))))"  # noqa: E501
+        "\n    values[1] = dx_dt"
+        "\nend"
+        "\n"
+    )