modelbase2 0.5.0__py3-none-any.whl → 0.7.0__py3-none-any.whl

modelbase2/__init__.py

@@ -50,6 +50,7 @@ __all__ = [
     "distributions",
     "experimental",
     "fit",
+    "fns",
     "make_protocol",
     "mc",
     "mca",
@@ -75,6 +76,7 @@ from . import (
     distributions,
     experimental,
     fit,
+    fns,
     mc,
     mca,
     nnarchitectures,

modelbase2/experimental/codegen.py

@@ -1,10 +1,10 @@
 """Module to export models as code."""
 
 import ast
-import inspect
 import warnings
 from collections.abc import Callable, Generator, Iterable, Iterator
 
+from modelbase2.experimental.source_tools import get_fn_ast, get_fn_source
 from modelbase2.model import Model
 from modelbase2.types import Derived
 
@@ -15,7 +15,6 @@ __all__ = [
     "conditional_join",
     "generate_model_code_py",
     "generate_modelbase_code",
-    "get_fn_source",
     "handle_fn",
 ]
 
@@ -55,28 +54,9 @@ class ReturnRemover(ast.NodeTransformer):
         return node.value
 
 
-def get_fn_source(fn: Callable) -> ast.FunctionDef:
-    """Get the source code of a function as an AST."""
-    import inspect
-    import textwrap
-
-    import dill
-
-    try:
-        source = inspect.getsource(fn)
-    except OSError:  # could not get source code
-        source = dill.source.getsource(fn)
-
-    tree = ast.parse(textwrap.dedent(source))
-    if not isinstance(fn_def := tree.body[0], ast.FunctionDef):
-        msg = "Not a function"
-        raise TypeError(msg)
-    return fn_def
-
-
 def handle_fn(fn: Callable, args: list[str]) -> str:
     """Get the source code of a function, removing docstrings and return statements."""
-    tree = get_fn_source(fn)
+    tree = get_fn_ast(fn)
 
     argmap = dict(zip([i.arg for i in tree.args.args], args, strict=True))
     tree = DocstringRemover().visit(tree)
@@ -118,31 +98,41 @@ def generate_modelbase_code(model: Model) -> str:
 
     # Derived
     derived_source = []
-    for k, v in model.derived.items():
-        fn = v.fn
+    for k, rxn in model.derived.items():
+        fn = rxn.fn
         fn_name = fn.__name__
-        functions[fn_name] = inspect.getsource(fn)
+        functions[fn_name] = get_fn_source(fn)
 
         derived_source.append(
-            f""".add_derived(
+            f"""        .add_derived(
                 "{k}",
                 fn={fn_name},
-                args={v.args},
+                args={rxn.args},
             )"""
         )
 
     # Reactions
     reactions_source = []
-    for k, v in model.reactions.items():
-        fn = v.fn
+    for k, rxn in model.reactions.items():
+        fn = rxn.fn
         fn_name = fn.__name__
-        functions[fn_name] = inspect.getsource(fn)
+        functions[fn_name] = get_fn_source(fn)
+        stoichiometry: list[str] = []
+        for var, stoich in rxn.stoichiometry.items():
+            if isinstance(stoich, Derived):
+                functions[fn_name] = get_fn_source(fn)
+                args = ", ".join(f'"{k}"' for k in stoich.args)
+                stoich = (  # noqa: PLW2901
+                    f"""Derived(name="{var}", fn={fn.__name__}, args=[{args}])"""
+                )
+            stoichiometry.append(f""""{var}": {stoich}""")
+
         reactions_source.append(
             f"""        .add_reaction(
                 "{k}",
                 fn={fn_name},
-                args={v.args},
-                stoichiometry={v.stoichiometry},
+                args={rxn.args},
+                stoichiometry={{{",".join(stoichiometry)}}},
             )"""
         )
 
@@ -178,25 +168,32 @@ def generate_modelbase_code(model: Model) -> str:
 
 def generate_model_code_py(model: Model) -> str:
     """Transform the model into a single function, inlining the function calls."""
+    source = [
+        "from collections.abc import Iterable\n",
+        "from modelbase2.types import Float\n",
+        "def model(t: Float, y: Float) -> Iterable[Float]:",
+    ]
+
     # Variables
     variables = model.variables
-    variable_source = "    {} = y".format(", ".join(variables))
+    if len(variables) > 0:
+        source.append("    {} = y".format(", ".join(variables)))
 
     # Parameters
-    parameter_source = "\n".join(f"    {k} = {v}" for k, v in model.parameters.items())
+    parameters = model.parameters
+    if len(parameters) > 0:
+        source.append("\n".join(f"    {k} = {v}" for k, v in model.parameters.items()))
 
     # Derived
-    derived_source = []
     for name, derived in model.derived.items():
-        derived_source.append(f"    {name} = {handle_fn(derived.fn, derived.args)}")
+        source.append(f"    {name} = {handle_fn(derived.fn, derived.args)}")
 
     # Reactions
-    reactions = []
     for name, rxn in model.reactions.items():
-        reactions.append(f"    {name} = {handle_fn(rxn.fn, rxn.args)}")
+        source.append(f"    {name} = {handle_fn(rxn.fn, rxn.args)}")
 
     # Stoichiometries
-    stoichiometries = {}
+    stoich_srcs = {}
     for rxn_name, rxn in model.reactions.items():
         for cpd_name, factor in rxn.stoichiometry.items():
             if isinstance(factor, Derived):
@@ -207,10 +204,9 @@ def generate_model_code_py(model: Model) -> str:
                 src = f"- {rxn_name}"
             else:
                 src = f"{factor} * {rxn_name}"
-            stoichiometries.setdefault(cpd_name, []).append(src)
-    stoich_source = []
-    for variable, stoich in stoichiometries.items():
-        stoich_source.append(
+            stoich_srcs.setdefault(cpd_name, []).append(src)
+    for variable, stoich in stoich_srcs.items():
+        source.append(
             f"    d{variable}dt = {conditional_join(stoich, lambda x: x.startswith('-'), ' ', ' + ')}"
         )
 
@@ -221,19 +217,14 @@ def generate_model_code_py(model: Model) -> str:
             stacklevel=1,
         )
 
-    # Combine all the sources
-    source = [
-        "from collections.abc import Iterable\n",
-        "from modelbase2.types import Float\n",
-        "def model(t: Float, y: Float) -> Iterable[Float]:",
-        variable_source,
-        parameter_source,
-        *derived_source,
-        *reactions,
-        *stoich_source,
-        "    return {}".format(
-            ", ".join(f"d{i}dt" for i in variables),
-        ),
-    ]
+    # Return
+    if len(variables) > 0:
+        source.append(
+            "    return {}".format(
+                ", ".join(f"d{i}dt" for i in variables),
+            ),
+        )
+    else:
+        source.append("    return ()")
 
     return "\n".join(source)

modelbase2/experimental/diff.py

@@ -1,7 +1,7 @@
 """Diffing utilities for comparing models."""
 
+from collections.abc import Mapping
 from dataclasses import dataclass, field
-from typing import Mapping
 
 from modelbase2.model import Model
 from modelbase2.types import Derived
@@ -128,9 +128,8 @@ def _soft_eq_stoichiometries(
                 return False
             if v1.args != v2.args:
                 return False
-        else:
-            if v1 != v2:
-                return False
+        elif v1 != v2:
+            return False
 
     return True
 

modelbase2/experimental/source_tools.py

@@ -0,0 +1,30 @@
+"""Tools for working with python source files."""
+
+import ast
+import inspect
+import textwrap
+from collections.abc import Callable
+
+import dill
+
+__all__ = [
+    "get_fn_ast",
+    "get_fn_source",
+]
+
+
+def get_fn_source(fn: Callable) -> str:
+    """Get the string representation of a function."""
+    try:
+        return inspect.getsource(fn)
+    except OSError:  # could not get source code
+        return dill.source.getsource(fn)
+
+
+def get_fn_ast(fn: Callable) -> ast.FunctionDef:
+    """Get the source code of a function as an AST."""
+    tree = ast.parse(textwrap.dedent(get_fn_source(fn)))
+    if not isinstance(fn_def := tree.body[0], ast.FunctionDef):
+        msg = "Not a function"
+        raise TypeError(msg)
+    return fn_def

modelbase2/experimental/symbolic.py

@@ -2,22 +2,43 @@
 
 import ast
 import inspect
-import textwrap
 from collections.abc import Callable
 from dataclasses import dataclass
+from types import ModuleType
 from typing import Any, cast
 
 import sympy
 
+from modelbase2.experimental.source_tools import get_fn_ast
 from modelbase2.model import Model
 
-__all__ = ["Context", "SymbolicModel", "model_fn_to_sympy", "to_symbolic_model"]
+__all__ = [
+    "Context",
+    "SymbolicModel",
+    "model_fn_to_sympy",
+    "to_symbolic_model",
+]
 
 
 @dataclass
 class Context:
     symbols: dict[str, sympy.Symbol | sympy.Expr]
     caller: Callable
+    parent_module: ModuleType | None
+
+    def updated(
+        self,
+        symbols: dict[str, sympy.Symbol | sympy.Expr] | None = None,
+        caller: Callable | None = None,
+        parent_module: ModuleType | None = None,
+    ) -> "Context":
+        return Context(
+            symbols=self.symbols if symbols is None else symbols,
+            caller=self.caller if caller is None else caller,
+            parent_module=self.parent_module
+            if parent_module is None
+            else parent_module,
+        )
 
 
 @dataclass
@@ -27,68 +48,21 @@ class SymbolicModel:
     eqs: list[sympy.Expr]
 
 
-def to_symbolic_model(model: Model) -> SymbolicModel:
-    cache = model._create_cache()  # noqa: SLF001
-
-    variables = dict(
-        zip(model.variables, sympy.symbols(list(model.variables)), strict=True)
-    )
-    parameters = dict(
-        zip(model.parameters, sympy.symbols(list(model.parameters)), strict=True)
-    )
-    symbols = variables | parameters
-
-    for k, v in model.derived.items():
-        symbols[k] = model_fn_to_sympy(v.fn, [symbols[i] for i in v.args])
-
-    rxns = {
-        k: model_fn_to_sympy(v.fn, [symbols[i] for i in v.args])
-        for k, v in model.reactions.items()
-    }
-
-    eqs: dict[str, sympy.Expr] = {}
-    for cpd, stoich in cache.stoich_by_cpds.items():
-        for rxn, stoich_value in stoich.items():
-            eqs[cpd] = (
-                eqs.get(cpd, sympy.Float(0.0)) + sympy.Float(stoich_value) * rxns[rxn]  # type: ignore
-            )
-
-    for cpd, dstoich in cache.dyn_stoich_by_cpds.items():
-        for rxn, der in dstoich.items():
-            eqs[cpd] = eqs.get(cpd, sympy.Float(0.0)) + model_fn_to_sympy(
-                der.fn,
-                [symbols[i] for i in der.args] * rxns[rxn],  # type: ignore
-            )  # type: ignore
-
-    return SymbolicModel(
-        variables=variables,
-        parameters=parameters,
-        eqs=[eqs[i] for i in cache.var_names],
-    )
-
-
 def model_fn_to_sympy(
     fn: Callable, model_args: list[sympy.Symbol | sympy.Expr] | None = None
 ) -> sympy.Expr:
-    source = textwrap.dedent(inspect.getsource(fn))
-
-    if not isinstance(fn_def := ast.parse(source).body[0], ast.FunctionDef):
-        msg = "Expected a function definition"
-        raise TypeError(msg)
-
+    fn_def = get_fn_ast(fn)
     fn_args = [str(arg.arg) for arg in fn_def.args.args]
-
     sympy_expr = _handle_fn_body(
         fn_def.body,
         ctx=Context(
             symbols={name: sympy.Symbol(name) for name in fn_args},
             caller=fn,
+            parent_module=inspect.getmodule(fn),
         ),
     )
-
     if model_args is not None:
         sympy_expr = sympy_expr.subs(dict(zip(fn_args, model_args, strict=True)))
-
     return cast(sympy.Expr, sympy_expr)
 
 
@@ -215,18 +189,45 @@ def _handle_binop(node: ast.BinOp, ctx: Context) -> sympy.Expr:
 
 
 def _handle_call(node: ast.Call, ctx: Context) -> sympy.Expr:
-    if not isinstance(callee := node.func, ast.Name):
-        msg = "Only function calls with names are supported"
-        raise TypeError(msg)
+    # direct call, e.g. mass_action(x, k1)
+    if isinstance(callee := node.func, ast.Name):
+        fn_name = str(callee.id)
+        fns = dict(inspect.getmembers(ctx.parent_module, predicate=callable))
+
+        return model_fn_to_sympy(
+            fns[fn_name],
+            model_args=[_handle_expr(i, ctx) for i in node.args],
+        )
+
+    # search for fn in other namespace
+    if isinstance(attr := node.func, ast.Attribute):
+        imports = dict(inspect.getmembers(ctx.parent_module, inspect.ismodule))
+
+        # Single level, e.g. fns.mass_action(x, k1)
+        if isinstance(module_name := attr.value, ast.Name):
+            return _handle_call(
+                ast.Call(func=ast.Name(attr.attr), args=node.args),
+                ctx=ctx.updated(parent_module=imports[module_name.id]),
+            )
 
-    fn_name = str(callee.id)
-    parent_module = inspect.getmodule(ctx.caller)
-    fns = dict(inspect.getmembers(parent_module, predicate=callable))
+        # Multiple levels, e.g. modelbase2.fns.mass_action(x, k1)
+        if isinstance(inner_attr := attr.value, ast.Attribute):
+            if not isinstance(module_name := inner_attr.value, ast.Name):
+                msg = f"Unknown target kind {module_name}"
+                raise NotImplementedError(msg)
+            return _handle_call(
+                ast.Call(
+                    func=ast.Attribute(
+                        value=ast.Name(inner_attr.attr),
+                        attr=attr.attr,
+                    ),
+                    args=node.args,
+                ),
+                ctx=ctx.updated(parent_module=imports[module_name.id]),
+            )
 
-    return model_fn_to_sympy(
-        fns[fn_name],
-        model_args=[_handle_expr(i, ctx) for i in node.args],
-    )
+    msg = f"Onsupported function type {node.func}"
+    raise NotImplementedError(msg)
 
 
 def _handle_name(node: ast.Name, ctx: Context) -> sympy.Symbol | sympy.Expr:
@@ -284,3 +285,43 @@ def _handle_expr(node: ast.expr, ctx: Context) -> sympy.Expr:
 
     msg = f"Expression type {type(node).__name__} not implemented"
     raise NotImplementedError(msg)
+
+
+def to_symbolic_model(model: Model) -> SymbolicModel:
+    cache = model._create_cache()  # noqa: SLF001
+
+    variables = dict(
+        zip(model.variables, sympy.symbols(list(model.variables)), strict=True)
+    )
+    parameters = dict(
+        zip(model.parameters, sympy.symbols(list(model.parameters)), strict=True)
+    )
+    symbols = variables | parameters
+
+    for k, v in model.derived.items():
+        symbols[k] = model_fn_to_sympy(v.fn, [symbols[i] for i in v.args])
+
+    rxns = {
+        k: model_fn_to_sympy(v.fn, [symbols[i] for i in v.args])
+        for k, v in model.reactions.items()
+    }
+
+    eqs: dict[str, sympy.Expr] = {}
+    for cpd, stoich in cache.stoich_by_cpds.items():
+        for rxn, stoich_value in stoich.items():
+            eqs[cpd] = (
+                eqs.get(cpd, sympy.Float(0.0)) + sympy.Float(stoich_value) * rxns[rxn]  # type: ignore
+            )
+
+    for cpd, dstoich in cache.dyn_stoich_by_cpds.items():
+        for rxn, der in dstoich.items():
+            eqs[cpd] = eqs.get(cpd, sympy.Float(0.0)) + model_fn_to_sympy(
+                der.fn,
+                [symbols[i] for i in der.args] * rxns[rxn],  # type: ignore
+            )  # type: ignore
+
+    return SymbolicModel(
+        variables=variables,
+        parameters=parameters,
+        eqs=[eqs[i] for i in cache.var_names],
+    )

modelbase2/experimental/tex.py

@@ -2,16 +2,21 @@
 
 from __future__ import annotations
 
-import ast
-import inspect
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, cast
 
 import latexify
 
+from modelbase2.experimental.source_tools import get_fn_ast
 from modelbase2.types import Derived, RateFn
 
-__all__ = ["TexExport", "TexReaction", "default_init", "get_model_tex_diff", "to_tex"]
+__all__ = [
+    "TexExport",
+    "TexReaction",
+    "default_init",
+    "get_model_tex_diff",
+    "to_tex",
+]
 
 if TYPE_CHECKING:
     from collections.abc import Callable, Mapping
@@ -67,9 +72,7 @@ def _escape_non_math(s: str) -> str:
 
 
 def _fn_to_latex(fn: Callable, arg_names: list[str]) -> str:
-    code = inspect.getsource(fn)
-    src = cast(ast.Module, ast.parse(code))
-    fn_def = cast(ast.FunctionDef, src.body[0])
+    fn_def = get_fn_ast(fn)
     args: list[str] = [i.arg for i in fn_def.args.args]
     arg_mapping: dict[str, str] = dict(zip(args, arg_names, strict=True))
     return cast(
@@ -329,7 +332,9 @@ class TexExport:
             parameters={gls.get(k, k): v for k, v in self.parameters.items()},
             variables={gls.get(k, k): v for k, v in self.variables.items()},
             derived={
-                gls.get(k, k): Derived(fn=v.fn, args=[gls.get(i, i) for i in v.args])
+                gls.get(k, k): Derived(
+                    name=k, fn=v.fn, args=[gls.get(i, i) for i in v.args]
+                )
                 for k, v in self.derived.items()
             },
             reactions={

modelbase2/fit.py

@@ -19,7 +19,14 @@ from scipy.optimize import minimize
 
 from modelbase2.integrators import DefaultIntegrator
 from modelbase2.simulator import Simulator
-from modelbase2.types import Array, ArrayLike, Callable, IntegratorProtocol, cast
+from modelbase2.types import (
+    Array,
+    ArrayLike,
+    Callable,
+    IntegratorProtocol,
+    cast,
+    unwrap,
+)
 
 __all__ = [
     "InitialGuess",
@@ -117,7 +124,7 @@ def _steady_state_residual(
         float: Root mean square error between model and data
 
     """
-    c_ss, v_ss = (
+    c_ss, v_ss = unwrap(
         Simulator(
             model.update_parameters(
                 dict(
@@ -132,7 +139,7 @@ def _steady_state_residual(
             integrator=integrator,
         )
         .simulate_to_steady_state()
-        .get_full_concs_and_fluxes()
+        .get_result()
     )
     if c_ss is None or v_ss is None:
         return cast(float, np.inf)
@@ -171,7 +178,7 @@ def _time_course_residual(
             integrator=integrator,
         )
         .simulate_time_course(data.index)  # type: ignore
-        .get_full_concs_and_fluxes()
+        .get_result()
     )
     if c_ss is None or v_ss is None:
         return cast(float, np.inf)

modelbase2/fns.py

@@ -8,6 +8,7 @@ if TYPE_CHECKING:
     from modelbase2.types import Float
 
 __all__ = [
+    "add",
     "constant",
     "diffusion_1s_1p",
     "div",
@@ -75,6 +76,11 @@ def twice(x: Float) -> Float:
     return x * 2
 
 
+def add(x: Float, y: Float) -> Float:
+    """Proportional function."""
+    return x + y
+
+
 def proportional(x: Float, y: Float) -> Float:
     """Proportional function."""
     return x * y

modelbase2/integrators/int_assimulo.py

@@ -22,7 +22,7 @@ with contextlib.redirect_stderr(open(os.devnull, "w")):  # noqa: PTH123
 if TYPE_CHECKING:
     from collections.abc import Callable
 
-    from modelbase2.types import ArrayLike
+    from modelbase2.types import Array, ArrayLike
 
 
 @dataclass
@@ -77,7 +77,7 @@ class Assimulo:
         *,
         t_end: float,
         steps: int | None = None,
-    ) -> tuple[ArrayLike | None, ArrayLike | None]:
+    ) -> tuple[Array | None, ArrayLike | None]:
         """Integrate the ODE system.
 
         Args:
@@ -100,7 +100,7 @@ class Assimulo:
         self,
         *,
         time_points: ArrayLike,
-    ) -> tuple[ArrayLike | None, ArrayLike | None]:
+    ) -> tuple[Array | None, ArrayLike | None]:
         """Integrate the ODE system over a time course.
 
         Args:

modelbase2/integrators/int_scipy.py

@@ -14,7 +14,7 @@ from typing import TYPE_CHECKING, cast
 import numpy as np
 import scipy.integrate as spi
 
-from modelbase2.types import ArrayLike, Float
+from modelbase2.types import Array, ArrayLike
 
 if TYPE_CHECKING:
     from collections.abc import Callable
@@ -66,7 +66,7 @@ class Scipy:
         *,
         t_end: float,
         steps: int | None = None,
-    ) -> tuple[ArrayLike | None, ArrayLike | None]:
+    ) -> tuple[Array | None, ArrayLike | None]:
         """Integrate the ODE system.
 
         Args:
@@ -87,7 +87,7 @@ class Scipy:
 
     def integrate_time_course(
         self, *, time_points: ArrayLike
-    ) -> tuple[ArrayLike | None, ArrayLike | None]:
+    ) -> tuple[Array | None, ArrayLike | None]:
         """Integrate the ODE system over a time course.
 
         Args:
@@ -97,7 +97,6 @@ class Scipy:
             tuple[ArrayLike, ArrayLike]: Tuple containing the time points and the integrated values.
 
         """
-
         y = spi.odeint(
             func=self.rhs,
             y0=self.y0,
@@ -108,7 +107,7 @@ class Scipy:
         )
         self.t0 = time_points[-1]
         self.y0 = y[-1, :]
-        return time_points, y
+        return np.array(time_points, dtype=float), y
 
     def integrate_to_steady_state(
         self,

modelbase2/linear_label_map.py

@@ -287,10 +287,12 @@ class LinearLabelMapper:
                 stoichiometry = {}
                 if substrate != "EXT":
                     stoichiometry[substrate] = Derived(
-                        _neg_one_div, [substrate.split("__")[0]]
+                        name=substrate, fn=_neg_one_div, args=[substrate.split("__")[0]]
                     )
                 if product != "EXT":
-                    stoichiometry[product] = Derived(_one_div, [product.split("__")[0]])
+                    stoichiometry[product] = Derived(
+                        name=product, fn=_one_div, args=[product.split("__")[0]]
+                    )
 
                 m.add_reaction(
                     name=f"{rxn_name}__{i}",