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

mxlpy/meta/codegen_mxlpy.py

@@ -3,94 +3,246 @@
 from __future__ import annotations
 
 import logging
-from typing import TYPE_CHECKING
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, cast
 
-from mxlpy.meta.sympy_tools import fn_to_sympy, list_of_symbols, sympy_to_python_fn
-from mxlpy.types import Derived
+import sympy
+from wadler_lindig import pformat
+
+from mxlpy.meta.sympy_tools import (
+    fn_to_sympy,
+    list_of_symbols,
+    sympy_to_inline_py,
+    sympy_to_python_fn,
+)
+from mxlpy.types import Derived, InitialAssignment
+from mxlpy.units import Quantity
 
 if TYPE_CHECKING:
-    import sympy
+    from collections.abc import Callable
 
     from mxlpy.model import Model
 
 __all__ = [
+    "SymbolicFn",
+    "SymbolicParameter",
+    "SymbolicReaction",
+    "SymbolicRepr",
+    "SymbolicVariable",
     "generate_mxlpy_code",
+    "generate_mxlpy_code_from_symbolic_repr",
 ]
 
 _LOGGER = logging.getLogger()
 
 
-def generate_mxlpy_code(model: Model) -> str:
-    """Generate a mxlpy model from a model."""
-    functions: dict[str, tuple[sympy.Expr, list[str]]] = {}
+@dataclass
+class SymbolicFn:
+    """Container for symbolic fn."""
 
-    # Variables and parameters
-    variables = model.get_raw_variables()
-    parameters = model.get_parameter_values()
+    fn_name: str
+    expr: sympy.Expr
+    args: list[str]
+
+    def __repr__(self) -> str:
+        """Return default representation."""
+        return pformat(self)
+
+
+@dataclass
+class SymbolicVariable:
+    """Container for symbolic variable."""
+
+    value: sympy.Float | SymbolicFn  # initial assignment
+    unit: Quantity | None
+
+    def __repr__(self) -> str:
+        """Return default representation."""
+        return pformat(self)
+
+
+@dataclass
+class SymbolicParameter:
+    """Container for symbolic par."""
+
+    value: sympy.Float | SymbolicFn  # initial assignment
+    unit: Quantity | None
+
+    def __repr__(self) -> str:
+        """Return default representation."""
+        return pformat(self)
+
+
+@dataclass
+class SymbolicReaction:
+    """Container for symbolic rxn."""
+
+    fn: SymbolicFn
+    stoichiometry: dict[str, sympy.Float | str | SymbolicFn]
+
+    def __repr__(self) -> str:
+        """Return default representation."""
+        return pformat(self)
+
+
+@dataclass
+class SymbolicRepr:
+    """Container for symbolic model."""
+
+    variables: dict[str, SymbolicVariable] = field(default_factory=dict)
+    parameters: dict[str, SymbolicParameter] = field(default_factory=dict)
+    derived: dict[str, SymbolicFn] = field(default_factory=dict)
+    reactions: dict[str, SymbolicReaction] = field(default_factory=dict)
+
+    def __repr__(self) -> str:
+        """Return default representation."""
+        return pformat(self)
+
+
+def _fn_to_symbolic_repr(k: str, fn: Callable, model_args: list[str]) -> SymbolicFn:
+    fn_name = fn.__name__
+    args = cast(list, list_of_symbols(model_args))
+    if (expr := fn_to_sympy(fn, origin=k, model_args=args)) is None:
+        msg = f"Unable to parse fn for '{k}'"
+        raise ValueError(msg)
+    return SymbolicFn(fn_name=fn_name, expr=expr, args=model_args)
+
+
+def _to_symbolic_repr(model: Model) -> SymbolicRepr:
+    sym = SymbolicRepr()
+
+    for k, variable in model.get_raw_variables().items():
+        sym.variables[k] = SymbolicVariable(
+            value=_fn_to_symbolic_repr(k, val.fn, val.args)
+            if isinstance(val := variable.initial_value, InitialAssignment)
+            else sympy.Float(val),
+            unit=cast(Quantity, variable.unit),
+        )
+
+    for k, parameter in model.get_raw_parameters().items():
+        sym.parameters[k] = SymbolicParameter(
+            value=_fn_to_symbolic_repr(k, val.fn, val.args)
+            if isinstance(val := parameter.value, InitialAssignment)
+            else sympy.Float(val),
+            unit=cast(Quantity, parameter.unit),
+        )
 
-    # Derived
-    derived_source = []
     for k, der in model.get_raw_derived().items():
-        fn = der.fn
-        fn_name = fn.__name__
-        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)
+        sym.derived[k] = _fn_to_symbolic_repr(k, der.fn, der.args)
 
-        functions[fn_name] = (expr, der.args)
+    for k, rxn in model.get_raw_reactions().items():
+        sym.reactions[k] = SymbolicReaction(
+            fn=_fn_to_symbolic_repr(k, rxn.fn, rxn.args),
+            stoichiometry={
+                k: _fn_to_symbolic_repr(k, v.fn, v.args)
+                if isinstance(v, Derived)
+                else sympy.Float(v)
+                for k, v in rxn.stoichiometry.items()
+            },
+        )
 
+    if len(model._surrogates) > 0:  # noqa: SLF001
+        msg = "Generating code for Surrogates not yet supported."
+        _LOGGER.warning(msg)
+    return sym
+
+
+def _codegen_variable(
+    k: str, var: SymbolicVariable, functions: dict[str, tuple[sympy.Expr, list[str]]]
+) -> str:
+    if isinstance(init := var.value, SymbolicFn):
+        fn_name = f"init_{init.fn_name}"
+        functions[fn_name] = (init.expr, init.args)
+        return f"""        .add_variable(
+            {k!r},
+            initial_value=InitialAssignment(fn={fn_name}, args={init.args!r}),
+        )"""
+
+    value = sympy_to_inline_py(init)
+    if (unit := var.unit) is not None:
+        return f"        .add_variable({k!r}, value={value}, unit={sympy_to_inline_py(unit)})"
+    return f"        .add_variable({k!r}, initial_value={value})"
+
+
+def _codegen_parameter(
+    k: str, par: SymbolicParameter, functions: dict[str, tuple[sympy.Expr, list[str]]]
+) -> str:
+    if isinstance(init := par.value, SymbolicFn):
+        fn_name = f"init_{init.fn_name}"
+        functions[fn_name] = (init.expr, init.args)
+        return f"""        .add_parameter(
+            {k!r},
+            value=InitialAssignment(fn={fn_name}, args={init.args!r}),
+        )"""
+
+    value = sympy_to_inline_py(init)
+    if (unit := par.unit) is not None:
+        return f"        .add_parameter({k!r}, value={value}, unit={sympy_to_inline_py(unit)})"
+    return f"        .add_parameter({k!r}, value={value})"
+
+
+def generate_mxlpy_code_from_symbolic_repr(
+    model: SymbolicRepr, imports: list[str] | None = None
+) -> str:
+    """Generate MxlPy source code from symbolic representation.
+
+    This is both used by MxlPy internally to codegen an existing model again and by the
+    SBML import to generate the file.
+    """
+    imports = [] if imports is None else imports
+
+    functions: dict[str, tuple[sympy.Expr, list[str]]] = {}
+
+    # Variables
+    variable_source = []
+    for k, var in model.variables.items():
+        variable_source.append(_codegen_variable(k, var, functions=functions))
+
+    # Parameters
+    parameter_source = []
+    for k, par in model.parameters.items():
+        parameter_source.append(_codegen_parameter(k, par, functions=functions))
+
+    # Derived
+    derived_source = []
+    for k, fn in model.derived.items():
+        functions[fn.fn_name] = (fn.expr, fn.args)
         derived_source.append(
             f"""        .add_derived(
-                "{k}",
-                fn={fn_name},
-                args={der.args},
+                {k!r},
+                fn={fn.fn_name},
+                args={fn.args},
             )"""
         )
 
     # Reactions
     reactions_source = []
-    for k, rxn in model.get_raw_reactions().items():
+    for k, rxn in model.reactions.items():
         fn = rxn.fn
-        fn_name = fn.__name__
-        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)
+        functions[fn.fn_name] = (fn.expr, fn.args)
+
         stoichiometry: list[str] = []
         for var, stoich in rxn.stoichiometry.items():
-            if isinstance(stoich, Derived):
-                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}])"""
+            if isinstance(stoich, SymbolicFn):
+                fn_name = f"{k}_stoich_{stoich.fn_name}"
+                functions[fn_name] = (stoich.expr, stoich.args)
+                stoichiometry.append(
+                    f""""{var}": Derived(fn={fn_name}, args={stoich.args!r})"""
                 )
-            stoichiometry.append(f""""{var}": {stoich}""")
-
+            elif isinstance(stoich, str):
+                stoichiometry.append(f""""{var}": {stoich!r}""")
+            else:
+                stoichiometry.append(f""""{var}": {sympy_to_inline_py(stoich)}""")
         reactions_source.append(
             f"""        .add_reaction(
                 "{k}",
-                fn={fn_name},
-                args={rxn.args},
+                fn={fn.fn_name},
+                args={fn.args},
                 stoichiometry={{{",".join(stoichiometry)}}},
             )"""
         )
 
     # Surrogates
-    if len(model._surrogates) > 0:  # noqa: SLF001
-        msg = "Generating code for Surrogates not yet supported."
-        _LOGGER.warning(msg)
 
     # Combine all the sources
     functions_source = "\n\n".join(
@@ -98,21 +250,26 @@ def generate_mxlpy_code(model: Model) -> str:
         for name, (expr, args) in functions.items()
     )
     source = [
-        "from mxlpy import Model\n",
+        *imports,
+        "from mxlpy import Model, Derived, InitialAssignment\n",
         functions_source,
+        "",
         "def create_model() -> Model:",
         "    return (",
         "        Model()",
     ]
-    if len(parameters) > 0:
-        source.append(f"        .add_parameters({parameters})")
-    if len(variables) > 0:
-        source.append(f"        .add_variables({variables})")
+    if len(variable_source) > 0:
+        source.append("\n".join(variable_source))
+    if len(parameter_source) > 0:
+        source.append("\n".join(parameter_source))
     if len(derived_source) > 0:
         source.append("\n".join(derived_source))
     if len(reactions_source) > 0:
         source.append("\n".join(reactions_source))
-
     source.append("    )")
-
     return "\n".join(source)
+
+
+def generate_mxlpy_code(model: Model) -> str:
+    """Generate a mxlpy model from a model."""
+    return generate_mxlpy_code_from_symbolic_repr(_to_symbolic_repr(model))

mxlpy/meta/source_tools.py

@@ -15,6 +15,7 @@ from typing import TYPE_CHECKING, Any, cast
 import dill
 import numpy as np
 import sympy
+from wadler_lindig import pformat
 
 if TYPE_CHECKING:
     from collections.abc import Callable
@@ -111,7 +112,6 @@ KNOWN_FNS: dict[Callable, sympy.Expr] = {
     math.trunc: sympy.trunc,
     # math.ulp: sympy.ulp,
     # numpy
-    np.exp: sympy.exp,
     np.abs: sympy.Abs,
     np.acos: sympy.acos,
     np.acosh: sympy.acosh,
@@ -175,6 +175,10 @@ class Context:
     modules: dict[str, ModuleType]
     fns: dict[str, Callable]
 
+    def __repr__(self) -> str:
+        """Return default representation."""
+        return pformat(self)
+
     def updated(
         self,
         symbols: dict[str, sympy.Symbol | sympy.Expr] | None = None,
@@ -309,10 +313,10 @@ def fn_to_sympy(
         )
         if sympy_expr is None:
             return None
-        # FIXME: we shouldn't end up here, where does this come from?
+        # Evaluated fns and floats from attributes
         if isinstance(sympy_expr, float):
             return sympy.Float(sympy_expr)
-        if model_args is not None:
+        if model_args is not None and len(model_args):
             sympy_expr = sympy_expr.subs(dict(zip(fn_args, model_args, strict=True)))
         return cast(sympy.Expr, sympy_expr)
 
@@ -323,77 +327,6 @@ def fn_to_sympy(
         return None
 
 
-def _handle_name(node: ast.Name, ctx: Context) -> sympy.Symbol | sympy.Expr:
-    value = ctx.symbols.get(node.id)
-    if value is None:
-        global_variables = dict(
-            inspect.getmembers(
-                ctx.parent_module,
-                predicate=lambda x: isinstance(x, float),
-            )
-        )
-        value = sympy.Float(global_variables[node.id])
-    return value
-
-
-def _handle_expr(node: ast.expr, ctx: Context) -> sympy.Expr | None:
-    if isinstance(node, float):
-        return sympy.Float(node)
-    if isinstance(node, ast.UnaryOp):
-        return _handle_unaryop(node, ctx)
-    if isinstance(node, ast.BinOp):
-        return _handle_binop(node, ctx)
-    if isinstance(node, ast.Name):
-        return _handle_name(node, ctx)
-    if isinstance(node, ast.Constant):
-        return node.value
-    if isinstance(node, ast.Call):
-        return _handle_call(node, ctx=ctx)
-    if isinstance(node, ast.Attribute):
-        return _handle_attribute(node, ctx=ctx)
-
-    if isinstance(node, ast.Compare):
-        # Handle chained comparisons like 1 < a < 2
-        left = cast(Any, _handle_expr(node.left, ctx))
-        comparisons = []
-
-        # Build all individual comparisons from the chain
-        prev_value = left
-        for op, comparator in zip(node.ops, node.comparators, strict=True):
-            right = cast(Any, _handle_expr(comparator, ctx))
-
-            if isinstance(op, ast.Gt):
-                comparisons.append(prev_value > right)
-            elif isinstance(op, ast.GtE):
-                comparisons.append(prev_value >= right)
-            elif isinstance(op, ast.Lt):
-                comparisons.append(prev_value < right)
-            elif isinstance(op, ast.LtE):
-                comparisons.append(prev_value <= right)
-            elif isinstance(op, ast.Eq):
-                comparisons.append(prev_value == right)
-            elif isinstance(op, ast.NotEq):
-                comparisons.append(prev_value != right)
-
-            prev_value = right
-
-        # Combine all comparisons with logical AND
-        result = comparisons[0]
-        for comp in comparisons[1:]:
-            result = sympy.And(result, comp)
-        return cast(sympy.Expr, result)
-
-    # Handle conditional expressions (ternary operators)
-    if isinstance(node, ast.IfExp):
-        condition = _handle_expr(node.test, ctx)
-        if_true = _handle_expr(node.body, ctx)
-        if_false = _handle_expr(node.orelse, ctx)
-        return sympy.Piecewise((if_true, condition), (if_false, True))
-
-    msg = f"Expression type {type(node).__name__} not implemented"
-    raise NotImplementedError(msg)
-
-
 def _handle_fn_body(body: list[ast.stmt], ctx: Context) -> sympy.Expr | None:
     pieces = []
     remaining_body = list(body)
@@ -505,6 +438,81 @@ def _handle_fn_body(body: list[ast.stmt], ctx: Context) -> sympy.Expr | None:
     raise ValueError(msg)
 
 
+def _handle_expr(node: ast.expr, ctx: Context) -> sympy.Expr | None:
+    """Key dispatch function."""
+    if isinstance(node, float):
+        return sympy.Float(node)
+    if isinstance(node, ast.UnaryOp):
+        return _handle_unaryop(node, ctx)
+    if isinstance(node, ast.BinOp):
+        return _handle_binop(node, ctx)
+    if isinstance(node, ast.Name):
+        return _handle_name(node, ctx)
+    if isinstance(node, ast.Constant):
+        if isinstance(val := node.value, (float, int)):
+            return sympy.Float(val)
+        msg = "Can only use float values"
+        raise NotImplementedError(msg)
+    if isinstance(node, ast.Call):
+        return _handle_call(node, ctx=ctx)
+    if isinstance(node, ast.Attribute):
+        return _handle_attribute(node, ctx=ctx)
+
+    if isinstance(node, ast.Compare):
+        # Handle chained comparisons like 1 < a < 2
+        left = cast(Any, _handle_expr(node.left, ctx))
+        comparisons = []
+
+        # Build all individual comparisons from the chain
+        prev_value = left
+        for op, comparator in zip(node.ops, node.comparators, strict=True):
+            right = cast(Any, _handle_expr(comparator, ctx))
+
+            if isinstance(op, ast.Gt):
+                comparisons.append(prev_value > right)
+            elif isinstance(op, ast.GtE):
+                comparisons.append(prev_value >= right)
+            elif isinstance(op, ast.Lt):
+                comparisons.append(prev_value < right)
+            elif isinstance(op, ast.LtE):
+                comparisons.append(prev_value <= right)
+            elif isinstance(op, ast.Eq):
+                comparisons.append(prev_value == right)
+            elif isinstance(op, ast.NotEq):
+                comparisons.append(prev_value != right)
+
+            prev_value = right
+
+        # Combine all comparisons with logical AND
+        result = comparisons[0]
+        for comp in comparisons[1:]:
+            result = sympy.And(result, comp)
+        return cast(sympy.Expr, result)
+
+    # Handle conditional expressions (ternary operators)
+    if isinstance(node, ast.IfExp):
+        condition = _handle_expr(node.test, ctx)
+        if_true = _handle_expr(node.body, ctx)
+        if_false = _handle_expr(node.orelse, ctx)
+        return sympy.Piecewise((if_true, condition), (if_false, True))
+
+    msg = f"Expression type {type(node).__name__} not implemented"
+    raise NotImplementedError(msg)
+
+
+def _handle_name(node: ast.Name, ctx: Context) -> sympy.Symbol | sympy.Expr:
+    value = ctx.symbols.get(node.id)
+    if value is None:
+        global_variables = dict(
+            inspect.getmembers(
+                ctx.parent_module,
+                predicate=lambda x: isinstance(x, float),
+            )
+        )
+        value = sympy.Float(global_variables[node.id])
+    return value
+
+
 def _handle_unaryop(node: ast.UnaryOp, ctx: Context) -> sympy.Expr:
     left = _handle_expr(node.operand, ctx)
     left = cast(Any, left)  # stupid sympy types don't allow ops on symbols
@@ -546,6 +554,27 @@ def _handle_binop(node: ast.BinOp, ctx: Context) -> sympy.Expr:
             raise NotImplementedError(msg)
 
 
+def _get_inner_object(obj: object, levels: list[str]) -> sympy.Float | None:
+    # Check if object is instantiated, otherwise instantiate first
+    if isinstance(obj, type):
+        obj = obj()
+
+    for level in levels:
+        _LOGGER.debug("obj %s, level %s", obj, level)
+        obj = getattr(obj, level, None)
+
+    if obj is None:
+        return None
+
+    if isinstance(obj, float):
+        if (value := KNOWN_CONSTANTS.get(obj)) is not None:
+            return value
+        return sympy.Float(obj)
+
+    _LOGGER.debug("Inner object not float: %s", obj)
+    return None
+
+
 # FIXME: check if target isn't an object or class
 def _handle_attribute(node: ast.Attribute, ctx: Context) -> sympy.Expr | None:
     """Handle an attribute.
@@ -581,17 +610,36 @@ def _handle_attribute(node: ast.Attribute, ctx: Context) -> sympy.Expr | None:
         dict(inspect.getmembers(ctx.parent_module, predicate=inspect.ismodule))
         | ctx.modules
     )
+    variables = vars(ctx.parent_module)
+
     match node.value:
         case ast.Name(l1):
             module_name = l1
             module = modules.get(module_name)
+            if module is None and (var := variables.get(l1)) is not None:
+                return _get_inner_object(var, [node.attr])
         case ast.Attribute():
-            levels = _find_root(node.value, [])
+            levels = _find_root(node.value, levels=[])
+            _LOGGER.debug("Attribute levels %s", levels)
             module_name = ".".join(levels)
-            for level in levels[:-1]:
-                modules.update(
-                    dict(inspect.getmembers(modules[level], predicate=inspect.ismodule))
-                )
+
+            for idx, level in enumerate(levels[:-1]):
+                if (module := modules.get(level)) is not None:
+                    modules.update(
+                        dict(
+                            inspect.getmembers(
+                                module,
+                                predicate=inspect.ismodule,
+                            )
+                        )
+                    )
+                elif (var := variables.get(level)) is not None:
+                    _LOGGER.debug("var %s", var)
+                    return _get_inner_object(var, levels[(idx + 1) :] + [node.attr])
+
+                else:
+                    _LOGGER.debug("No target found")
+
             module = modules.get(levels[-1])
         case _:
             raise NotImplementedError
@@ -633,6 +681,7 @@ def _handle_call(node: ast.Call, ctx: Context) -> sympy.Expr | None:
         if (expr := _handle_expr(i, ctx)) is None:
             return None
         model_args.append(expr)
+    _LOGGER.debug("Fn args: %s", model_args)
 
     match node.func:
         case ast.Name(id):
@@ -674,7 +723,7 @@ def _handle_call(node: ast.Call, ctx: Context) -> sympy.Expr | None:
         return None
 
     if (fn := KNOWN_FNS.get(py_fn)) is not None:
-        return fn
+        return sympy.Float(fn(*model_args))  # type: ignore
 
     return fn_to_sympy(
         py_fn,

mxlpy/meta/sympy_tools.py

@@ -5,7 +5,7 @@ from __future__ import annotations
 from typing import TYPE_CHECKING, cast
 
 import sympy
-from sympy.printing import rust_code
+from sympy.printing import jscode, rust_code
 from sympy.printing.pycode import pycode
 
 from mxlpy.meta.source_tools import fn_to_sympy
@@ -17,6 +17,7 @@ if TYPE_CHECKING:
 __all__ = [
     "list_of_symbols",
     "stoichiometries_to_sympy",
+    "sympy_to_inline_js",
     "sympy_to_inline_py",
     "sympy_to_inline_rust",
     "sympy_to_python_fn",
@@ -50,12 +51,17 @@ def sympy_to_inline_py(expr: sympy.Expr) -> str:
     'x**2 + 2*x + 1'
 
     """
-    return cast(str, pycode(expr, fully_qualified_modules=True))
+    return cast(str, pycode(expr, fully_qualified_modules=True, full_prec=False))
+
+
+def sympy_to_inline_js(expr: sympy.Expr) -> str:
+    """Create rust code from sympy expression."""
+    return cast(str, jscode(expr, full_prec=False))
 
 
 def sympy_to_inline_rust(expr: sympy.Expr) -> str:
     """Create rust code from sympy expression."""
-    return cast(str, rust_code(expr))
+    return cast(str, rust_code(expr, full_prec=False))
 
 
 def sympy_to_python_fn(
@@ -93,8 +99,8 @@ def sympy_to_python_fn(
     fn_args = ", ".join(f"{i}: float" for i in args)
 
     return f"""def {fn_name}({fn_args}) -> float:
-    return {pycode(expr)}
-    """
+    return {pycode(expr, fully_qualified_modules=True, full_prec=False)}
+    """.replace("math.factorial", "scipy.special.factorial")
 
 
 def stoichiometries_to_sympy(
@@ -114,4 +120,4 @@ def stoichiometries_to_sympy(
             expr = expr + sympy_fn * sympy.Symbol(rxn_name)  # type: ignore
         else:
             expr = expr + rxn_stoich * sympy.Symbol(rxn_name)  # type: ignore
-    return expr
+    return expr.subs(1.0, 1)  # type: ignore