mxlpy 0.20.0__py3-none-any.whl → 0.22.0__py3-none-any.whl

mxlpy/meta/codegen_model.py

@@ -0,0 +1,174 @@
+"""Module to export models as code."""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING
+
+from mxlpy.meta.sympy_tools import (
+    fn_to_sympy,
+    list_of_symbols,
+    stoichiometries_to_sympy,
+    sympy_to_inline_py,
+    sympy_to_inline_rust,
+)
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    import sympy
+
+    from mxlpy.model import Model
+
+__all__ = [
+    "generate_model_code_py",
+    "generate_model_code_rs",
+]
+
+_LOGGER = logging.getLogger(__name__)
+
+
+def _generate_model_code(
+    model: Model,
+    *,
+    sized: bool,
+    model_fn: str,
+    variables_template: str,
+    assignment_template: str,
+    sympy_inline_fn: Callable[[sympy.Expr], str],
+    return_template: str,
+    imports: list[str] | None = None,
+    end: str | None = None,
+    free_parameters: list[str] | None = None,
+) -> str:
+    source: list[str] = []
+    # Model components
+    variables = model.get_initial_conditions()
+    parameters = model.get_parameter_values()
+
+    if imports is not None:
+        source.extend(imports)
+
+    if not sized:
+        source.append(model_fn)
+    else:
+        source.append(model_fn.format(n=len(variables)))
+
+    if len(variables) > 0:
+        source.append(variables_template.format(", ".join(variables)))
+
+    # Parameters
+    if free_parameters is not None:
+        for key in free_parameters:
+            parameters.pop(key)
+    if len(parameters) > 0:
+        source.append(
+            "\n".join(
+                assignment_template.format(k=k, v=v) for k, v in parameters.items()
+            )
+        )
+
+    # Derived
+    for name, derived in model.get_raw_derived().items():
+        expr = fn_to_sympy(
+            derived.fn,
+            origin=name,
+            model_args=list_of_symbols(derived.args),
+        )
+        if expr is None:
+            msg = f"Unable to parse fn for derived value '{name}'"
+            raise ValueError(msg)
+        source.append(assignment_template.format(k=name, v=sympy_inline_fn(expr)))
+
+    # Reactions
+    for name, rxn in model.get_raw_reactions().items():
+        expr = fn_to_sympy(
+            rxn.fn,
+            origin=name,
+            model_args=list_of_symbols(rxn.args),
+        )
+        if expr is None:
+            msg = f"Unable to parse fn for reaction value '{name}'"
+            raise ValueError(msg)
+        source.append(assignment_template.format(k=name, v=sympy_inline_fn(expr)))
+
+    # Diff eqs
+    diff_eqs = {}
+    for rxn_name, rxn in model.get_raw_reactions().items():
+        for var_name, factor in rxn.stoichiometry.items():
+            diff_eqs.setdefault(var_name, {})[rxn_name] = factor
+
+    for variable, stoich in diff_eqs.items():
+        expr = stoichiometries_to_sympy(origin=variable, stoichs=stoich)
+        source.append(
+            assignment_template.format(k=f"d{variable}dt", v=sympy_inline_fn(expr))
+        )
+
+    # Surrogates
+    if len(model._surrogates) > 0:  # noqa: SLF001
+        msg = "Generating code for Surrogates not yet supported."
+        _LOGGER.warning(msg)
+
+    # Return
+    ret = ", ".join(f"d{i}dt" for i in diff_eqs) if len(diff_eqs) > 0 else "()"
+    source.append(return_template.format(ret))
+
+    if end is not None:
+        source.append(end)
+
+    # print(source)
+    return "\n".join(source)
+
+
+def generate_model_code_py(
+    model: Model,
+    free_parameters: list[str] | None = None,
+) -> str:
+    """Transform the model into a python function, inlining the function calls."""
+    if free_parameters is None:
+        model_fn = (
+            "def model(time: float, variables: Iterable[float]) -> Iterable[float]:"
+        )
+    else:
+        args = ", ".join(f"{k}: float" for k in free_parameters)
+        model_fn = f"def model(time: float, variables: Iterable[float], {args}) -> Iterable[float]:"
+
+    return _generate_model_code(
+        model,
+        imports=[
+            "from collections.abc import Iterable\n",
+        ],
+        sized=False,
+        model_fn=model_fn,
+        variables_template="    {} = variables",
+        assignment_template="    {k} = {v}",
+        sympy_inline_fn=sympy_to_inline_py,
+        return_template="    return {}",
+        end=None,
+        free_parameters=free_parameters,
+    )
+
+
+def generate_model_code_rs(
+    model: Model,
+    free_parameters: list[str] | None = None,
+) -> str:
+    """Transform the model into a rust function, inlining the function calls."""
+    if free_parameters is None:
+        model_fn = "fn model(time: f64, variables: &[f64; {n}]) -> [f64; {n}] {{"
+    else:
+        args = ", ".join(f"{k}: f64" for k in free_parameters)
+        model_fn = f"fn model(time: f64, variables: &[f64; {{n}}], {args}) -> [f64; {{n}}] {{{{"
+
+    return _generate_model_code(
+        model,
+        imports=None,
+        sized=True,
+        model_fn=model_fn,
+        variables_template="    let [{}] = *variables;",
+        assignment_template="    let {k}: f64 = {v};",
+        sympy_inline_fn=sympy_to_inline_rust,
+        return_template="    return [{}]",
+        end="}",
+        free_parameters=free_parameters,
+    )

mxlpy/meta/{codegen_modebase.py → codegen_mxlpy.py}

@@ -2,43 +2,44 @@
 
 from __future__ import annotations
 
-import warnings
+import logging
 from typing import TYPE_CHECKING
 
-import sympy
-
-from mxlpy.meta.source_tools import fn_to_sympy, sympy_to_fn
+from mxlpy.meta.sympy_tools import fn_to_sympy, list_of_symbols, sympy_to_python_fn
 from mxlpy.types import Derived
 
 if TYPE_CHECKING:
+    import sympy
+
     from mxlpy.model import Model
 
 __all__ = [
     "generate_mxlpy_code",
 ]
 
-
-def _list_of_symbols(args: list[str]) -> list[sympy.Symbol | sympy.Expr]:
-    return [sympy.Symbol(arg) for arg in args]
+_LOGGER = logging.getLogger()
 
 
 def generate_mxlpy_code(model: Model) -> str:
     """Generate a mxlpy model from a model."""
-    functions = {}
+    functions: dict[str, tuple[sympy.Expr, list[str]]] = {}
 
     # Variables and parameters
-    variables = model.variables
-    parameters = model.parameters
+    variables = model.get_raw_variables()
+    parameters = model.get_parameter_values()
 
     # Derived
     derived_source = []
-    for k, der in model.derived.items():
+    for k, der in model.get_raw_derived().items():
         fn = der.fn
         fn_name = fn.__name__
-        functions[fn_name] = (
-            fn_to_sympy(fn, model_args=_list_of_symbols(der.args)),
-            der.args,
-        )
+        if (
+            expr := fn_to_sympy(fn, origin=k, model_args=list_of_symbols(der.args))
+        ) is None:
+            msg = f"Unable to parse fn for derived value '{k}'"
+            raise ValueError(msg)
+
+        functions[fn_name] = (expr, der.args)
 
         derived_source.append(
             f"""        .add_derived(
@@ -50,20 +51,27 @@ def generate_mxlpy_code(model: Model) -> str:
 
     # Reactions
     reactions_source = []
-    for k, rxn in model.reactions.items():
+    for k, rxn in model.get_raw_reactions().items():
         fn = rxn.fn
         fn_name = fn.__name__
-        functions[fn_name] = (
-            fn_to_sympy(fn, model_args=_list_of_symbols(rxn.args)),
-            rxn.args,
-        )
+        if (
+            expr := fn_to_sympy(fn, origin=k, model_args=list_of_symbols(rxn.args))
+        ) is None:
+            msg = f"Unable to parse fn for reaction '{k}'"
+            raise ValueError(msg)
+
+        functions[fn_name] = (expr, rxn.args)
         stoichiometry: list[str] = []
         for var, stoich in rxn.stoichiometry.items():
             if isinstance(stoich, Derived):
-                functions[fn_name] = (
-                    fn_to_sympy(fn, model_args=_list_of_symbols(stoich.args)),
-                    rxn.args,
-                )
+                if (
+                    expr := fn_to_sympy(
+                        fn, origin=var, model_args=list_of_symbols(stoich.args)
+                    )
+                ) is None:
+                    msg = f"Unable to parse fn for stoichiometry '{var}'"
+                    raise ValueError(msg)
+                functions[fn_name] = (expr, rxn.args)
                 args = ", ".join(f'"{k}"' for k in stoich.args)
                 stoich = (  # noqa: PLW2901
                     f"""Derived(fn={fn.__name__}, args=[{args}])"""
@@ -81,14 +89,12 @@ def generate_mxlpy_code(model: Model) -> str:
 
     # Surrogates
     if len(model._surrogates) > 0:  # noqa: SLF001
-        warnings.warn(
-            "Generating code for Surrogates not yet supported.",
-            stacklevel=1,
-        )
+        msg = "Generating code for Surrogates not yet supported."
+        _LOGGER.warning(msg)
 
     # Combine all the sources
     functions_source = "\n\n".join(
-        sympy_to_fn(fn_name=name, args=args, expr=expr)
+        sympy_to_python_fn(fn_name=name, args=args, expr=expr)
         for name, (expr, args) in functions.items()
     )
     source = [