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

mxlpy/mca.py

@@ -71,7 +71,7 @@ def _response_coefficient_worker(
             - Series of flux response coefficients
 
     """
-    old = model.parameters[parameter]
+    old = model.get_parameter_values()[parameter]
     if y0 is not None:
         model.update_variables(y0)
 
@@ -205,7 +205,7 @@ def parameter_elasticities(
 
     variables = model.get_initial_conditions() if variables is None else variables
     for par in to_scan:
-        old = model.parameters[par]
+        old = model.get_parameter_values()[par]
 
         model.update_parameters({par: old * (1 + displacement)})
         upper = model.get_fluxes(variables=variables, time=time)

mxlpy/meta/__init__.py

@@ -2,12 +2,14 @@
 
 from __future__ import annotations
 
-from .codegen_latex import generate_latex_code
-from .codegen_modebase import generate_mxlpy_code
-from .codegen_py import generate_model_code_py
+from .codegen_latex import generate_latex_code, to_tex_export
+from .codegen_model import generate_model_code_py, generate_model_code_rs
+from .codegen_mxlpy import generate_mxlpy_code
 
 __all__ = [
     "generate_latex_code",
     "generate_model_code_py",
+    "generate_model_code_rs",
     "generate_mxlpy_code",
+    "to_tex_export",
 ]

mxlpy/meta/codegen_latex.py

@@ -7,7 +7,7 @@ from typing import TYPE_CHECKING
 
 import sympy
 
-from mxlpy.meta.source_tools import fn_to_sympy
+from mxlpy.meta.sympy_tools import fn_to_sympy, list_of_symbols
 from mxlpy.types import Derived, RateFn
 
 if TYPE_CHECKING:
@@ -21,6 +21,7 @@ __all__ = [
     "default_init",
     "generate_latex_code",
     "get_model_tex_diff",
+    "to_tex_export",
 ]
 
 cdot = r"\cdot"
@@ -31,10 +32,6 @@ newline = r"\\" + "\n"
 floatbarrier = r"\FloatBarrier"
 
 
-def _list_of_symbols(args: list[str]) -> list[sympy.Symbol | sympy.Expr]:
-    return [sympy.Symbol(arg) for arg in args]
-
-
 def default_init[T1, T2](d: dict[T1, T2] | None) -> dict[T1, T2]:
     """Return empty dict if d is None.
 
@@ -63,10 +60,6 @@ def _gls(s: str) -> str:
     return rf"\gls{{{s}}}"
 
 
-def _abbrev_and_full(s: str) -> str:
-    return rf"\acrfull{{{s}}}"
-
-
 def _gls_short(s: str) -> str:
     return rf"\acrshort{{{s}}}"
 
@@ -75,6 +68,10 @@ def _gls_full(s: str) -> str:
     return rf"\acrlong{{{s}}}"
 
 
+def _gls_short_and_full(s: str) -> str:
+    return rf"\acrfull{{{s}}}"
+
+
 def _rename_latex(s: str) -> str:
     if s[0].isdigit():
         s = s[1:]
@@ -109,6 +106,8 @@ def _sympy_to_latex(expr: sympy.Expr) -> str:
 
 def _fn_to_latex(
     fn: Callable,
+    *,
+    origin: str,
     arg_names: list[str],
     long_name_cutoff: int,
 ) -> tuple[str, dict[str, str]]:
@@ -121,10 +120,13 @@ def _fn_to_latex(
     replacements = {k: _name_to_latex(f"_x{i}") for i, k in enumerate(long_names)}
 
     expr = fn_to_sympy(
-        fn, _list_of_symbols([replacements.get(k, k) for k in tex_names])
+        fn,
+        origin=origin,
+        model_args=list_of_symbols([replacements.get(k, k) for k in tex_names]),
     )
-    fn_str = _sympy_to_latex(expr)
-    return fn_str, replacements
+    if expr is None:
+        return rf"\textcolor{{red}}{{{origin}}}", replacements
+    return _sympy_to_latex(expr), replacements
 
 
 def _table(
@@ -323,7 +325,8 @@ def _stoichs_to_latex(
             )
             sympy_fn = fn_to_sympy(
                 rxn_stoich.fn,
-                _list_of_symbols([replacements.get(k, k) for k in arg_names]),
+                origin=rxn_name,
+                model_args=list_of_symbols([replacements.get(k, k) for k in arg_names]),
             )
             expr = expr + sympy_fn * sympy.Symbol(rxn_name)  # type: ignore
         else:
@@ -480,7 +483,7 @@ class TexExport:
 
         def _add_gls_if_found(k: str) -> str:
             if (new := gls.get(k)) is not None:
-                return _abbrev_and_full(new)
+                return _gls_short_and_full(new)
             return k
 
         return TexExport(
@@ -564,7 +567,7 @@ class TexExport:
 
     def export_derived(
         self,
-        long_name_cutoff: int,
+        long_name_cutoff: int = 10,
     ) -> str:
         """Export derived quantities as LaTeX equations.
 
@@ -587,16 +590,20 @@ class TexExport:
         for k, v in sorted(self.derived.items()):
             fn_str, repls = _fn_to_latex(
                 v.fn,
+                origin=k,
                 arg_names=v.args,
                 long_name_cutoff=long_name_cutoff,
             )
-            rows.append(f"{_mathrm(_name_to_latex(k))} &= {fn_str} \\\\")
+            rows.append(f"    {_mathrm(_name_to_latex(k))} &= {fn_str} \\\\")
             if repls:
                 rows.append(_replacements_in_align(repls))
 
         return _latex_align(rows)
 
-    def export_reactions(self, long_name_cutoff: int) -> str:
+    def export_reactions(
+        self,
+        long_name_cutoff: int = 10,
+    ) -> str:
         """Export reactions as LaTeX equations.
 
         Returns
@@ -618,17 +625,18 @@ class TexExport:
         for k, v in sorted(self.reactions.items()):
             fn_str, repls = _fn_to_latex(
                 v.fn,
+                origin=k,
                 arg_names=v.args,
                 long_name_cutoff=long_name_cutoff,
             )
-            rows.append(f"{_mathrm(_name_to_latex(k))} &= {fn_str} \\\\")
+            rows.append(f"    {_mathrm(_name_to_latex(k))} &= {fn_str} \\\\")
             if repls:
                 rows.append(_replacements_in_align(repls))
         return _latex_align(rows)
 
     def export_diff_eqs(
         self,
-        long_name_cutoff: int,
+        long_name_cutoff: int = 10,
     ) -> str:
         """Export stoichiometries as LaTeX table.
 
@@ -654,12 +662,15 @@ class TexExport:
                 long_name_cutoff=long_name_cutoff,
             )
 
-            rows.append(f"{dxdt} &= {stoich_str} \\\\")
+            rows.append(f"    {dxdt} &= {stoich_str} \\\\")
             if repls:
                 rows.append(_replacements_in_align(repls))
         return _latex_align(rows)
 
-    def export_all(self, long_name_cutoff: int = 10) -> str:
+    def export_all(
+        self,
+        long_name_cutoff: int = 10,
+    ) -> str:
         """Export all model parts as a complete LaTeX document section.
 
         Returns
@@ -754,7 +765,7 @@ class TexExport:
 \usepackage[a4paper,top=2cm,bottom=2cm,left=2cm,right=2cm,marginparwidth=1.75cm]{{geometry}}
 \usepackage{{amsmath, amssymb, array, booktabs,
             breqn, caption, longtable, mathtools, placeins,
-            ragged2e, tabularx, titlesec, titling}}
+            ragged2e, tabularx, titlesec, titling, xcolor}}
 \newcommand{{\sectionbreak}}{{\clearpage}}
 \setlength{{\parindent}}{{0pt}}
 \allowdisplaybreaks
@@ -769,17 +780,20 @@ class TexExport:
 """
 
 
-def _to_tex_export(self: Model) -> TexExport:
+def to_tex_export(model: Model) -> TexExport:
+    """Create TexExport object from a model."""
     diff_eqs = {}
-    for rxn_name, rxn in self.reactions.items():
+    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
 
     return TexExport(
-        parameters=self.parameters,
-        variables=self.get_initial_conditions(),  # FIXME: think about this later
-        derived=self.derived,
-        reactions={k: TexReaction(v.fn, v.args) for k, v in self.reactions.items()},
+        parameters=model.get_parameter_values(),
+        variables=model.get_initial_conditions(),  # FIXME: think about this later
+        derived=model.get_raw_derived(),
+        reactions={
+            k: TexReaction(v.fn, v.args) for k, v in model.get_raw_reactions().items()
+        },
         diff_eqs=diff_eqs,
     )
 
@@ -820,7 +834,7 @@ def generate_latex_code(
     """
     gls = default_init(gls)
     return (
-        _to_tex_export(model)
+        to_tex_export(model)
         .rename_with_glossary(gls)
         .export_document(long_name_cutoff=long_name_cutoff)
     )
@@ -866,7 +880,7 @@ def get_model_tex_diff(
     return f"""{" start autogenerated ":%^60}
 {_clearpage()}
 {_subsubsection("Model changes")}{_label(section_label)}
-{((_to_tex_export(m1) - _to_tex_export(m2)).rename_with_glossary(gls).export_all(long_name_cutoff=long_name_cutoff))}
+{((to_tex_export(m1) - to_tex_export(m2)).rename_with_glossary(gls).export_all(long_name_cutoff=long_name_cutoff))}
 {_clearpage()}
 {" end autogenerated ":%^60}
 """

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 = [