mxlpy 0.15.0__py3-none-any.whl → 0.17.0__py3-none-any.whl

mxlpy/meta/source_tools.py

@@ -24,6 +24,8 @@ __all__ = [
 
 @dataclass
 class Context:
+    """Context for converting a function to sympy expression."""
+
     symbols: dict[str, sympy.Symbol | sympy.Expr]
     caller: Callable
     parent_module: ModuleType | None
@@ -34,6 +36,7 @@ class Context:
         caller: Callable | None = None,
         parent_module: ModuleType | None = None,
     ) -> "Context":
+        """Update the context with new values."""
         return Context(
             symbols=self.symbols if symbols is None else symbols,
             caller=self.caller if caller is None else caller,
@@ -44,7 +47,26 @@ class Context:
 
 
 def get_fn_source(fn: Callable) -> str:
-    """Get the string representation of a function."""
+    """Get the string representation of a function.
+
+    Parameters
+    ----------
+    fn
+        The function to extract source from
+
+    Returns
+    -------
+    str
+        String representation of the function's source code
+
+    Examples
+    --------
+    >>> def example_fn(x): return x * 2
+    >>> source = get_fn_source(example_fn)
+    >>> print(source)
+    def example_fn(x): return x * 2
+
+    """
     try:
         return inspect.getsource(fn)
     except OSError:  # could not get source code
@@ -52,7 +74,31 @@ def get_fn_source(fn: Callable) -> str:
 
 
 def get_fn_ast(fn: Callable) -> ast.FunctionDef:
-    """Get the source code of a function as an AST."""
+    """Get the source code of a function as an AST.
+
+    Parameters
+    ----------
+    fn
+        The function to convert to AST
+
+    Returns
+    -------
+    ast.FunctionDef
+        Abstract syntax tree representation of the function
+
+    Raises
+    ------
+    TypeError
+        If the input is not a function
+
+    Examples
+    --------
+    >>> def example_fn(x): return x * 2
+    >>> ast_tree = get_fn_ast(example_fn)
+    >>> isinstance(ast_tree, ast.FunctionDef)
+    True
+
+    """
     tree = ast.parse(textwrap.dedent(get_fn_source(fn)))
     if not isinstance(fn_def := tree.body[0], ast.FunctionDef):
         msg = "Not a function"
@@ -61,6 +107,27 @@ def get_fn_ast(fn: Callable) -> ast.FunctionDef:
 
 
 def sympy_to_inline(expr: sympy.Expr) -> str:
+    """Convert a sympy expression to inline Python code.
+
+    Parameters
+    ----------
+    expr
+        The sympy expression to convert
+
+    Returns
+    -------
+    str
+        Python code string for the expression
+
+    Examples
+    --------
+    >>> import sympy
+    >>> x = sympy.Symbol('x')
+    >>> expr = x**2 + 2*x + 1
+    >>> sympy_to_inline(expr)
+    'x**2 + 2*x + 1'
+
+    """
     return cast(str, pycode(expr, fully_qualified_modules=True))
 
 
@@ -70,7 +137,32 @@ def sympy_to_fn(
     args: list[str],
     expr: sympy.Expr,
 ) -> str:
-    """Convert a sympy expression to a python function."""
+    """Convert a sympy expression to a python function.
+
+    Parameters
+    ----------
+    fn_name
+        Name of the function to generate
+    args
+        List of argument names for the function
+    expr
+        Sympy expression to convert to a function body
+
+    Returns
+    -------
+    str
+        String representation of the generated function
+
+    Examples
+    --------
+    >>> import sympy
+    >>> x, y = sympy.symbols('x y')
+    >>> expr = x**2 + y
+    >>> print(sympy_to_fn(fn_name="square_plus_y", args=["x", "y"], expr=expr))
+    def square_plus_y(x: float, y: float) -> float:
+        return x**2 + y
+
+    """
     fn_args = ", ".join(f"{i}: float" for i in args)
 
     return f"""def {fn_name}({fn_args}) -> float:
@@ -82,7 +174,33 @@ def fn_to_sympy(
     fn: Callable,
     model_args: list[sympy.Symbol | sympy.Expr] | None = None,
 ) -> sympy.Expr:
-    """Convert a python function to a sympy expression."""
+    """Convert a python function to a sympy expression.
+
+    Parameters
+    ----------
+    fn
+        The function to convert
+    model_args
+        Optional list of sympy symbols to substitute for function arguments
+
+    Returns
+    -------
+    sympy.Expr
+        Sympy expression equivalent to the function
+
+    Examples
+    --------
+    >>> def square_fn(x):
+    ...     return x**2
+    >>> import sympy
+    >>> fn_to_sympy(square_fn)
+    x**2
+    >>> # With model_args
+    >>> y = sympy.Symbol('y')
+    >>> fn_to_sympy(square_fn, [y])
+    y**2
+
+    """
     fn_def = get_fn_ast(fn)
     fn_args = [str(arg.arg) for arg in fn_def.args.args]
     sympy_expr = _handle_fn_body(

mxlpy/model.py

@@ -18,6 +18,7 @@ import pandas as pd
 
 from mxlpy import fns
 from mxlpy.types import (
+    AbstractSurrogate,
     Array,
     Derived,
     Reaction,
@@ -27,6 +28,7 @@ from mxlpy.types import (
 __all__ = [
     "ArityMismatchError",
     "CircularDependencyError",
+    "Dependency",
     "MissingDependenciesError",
     "Model",
     "ModelCache",
@@ -36,7 +38,16 @@ if TYPE_CHECKING:
     from collections.abc import Iterable, Mapping
     from inspect import FullArgSpec
 
-    from mxlpy.types import AbstractSurrogate, Callable, Param, RateFn, RetType
+    from mxlpy.types import Callable, Param, RateFn, RetType
+
+
+@dataclass
+class Dependency:
+    """Container class for building dependency tree."""
+
+    name: str
+    required: set[str]
+    provided: set[str]
 
 
 class MissingDependenciesError(Exception):
@@ -145,30 +156,33 @@ def _invalidate_cache(method: Callable[Param, RetType]) -> Callable[Param, RetTy
 
 def _check_if_is_sortable(
     available: set[str],
-    elements: list[tuple[str, set[str]]],
+    elements: list[Dependency],
 ) -> None:
     all_available = available.copy()
-    for name, _ in elements:
-        all_available.add(name)
+    for dependency in elements:
+        all_available.update(dependency.provided)
 
     # Check if it can be sorted in the first place
     not_solvable = {}
-    for name, args in elements:
-        if not args.issubset(all_available):
-            not_solvable[name] = sorted(args.difference(all_available))
+    for dependency in elements:
+        if not dependency.required.issubset(all_available):
+            not_solvable[dependency.name] = sorted(
+                dependency.required.difference(all_available)
+            )
 
     if not_solvable:
         raise MissingDependenciesError(not_solvable=not_solvable)
 
 
 def _sort_dependencies(
-    available: set[str], elements: list[tuple[str, set[str]]]
+    available: set[str],
+    elements: list[Dependency],
 ) -> list[str]:
     """Sort model elements topologically based on their dependencies.
 
     Args:
         available: Set of available component names
-        elements: List of (name, dependencies) tuples to sort
+        elements: List of (name, dependencies, supplier) tuples to sort
 
     Returns:
         List of element names in dependency order
@@ -184,26 +198,27 @@ def _sort_dependencies(
     order = []
     # FIXME: what is the worst case here?
     max_iterations = len(elements) ** 2
-    queue: SimpleQueue[tuple[str, set[str]]] = SimpleQueue()
-    for k, v in elements:
-        queue.put((k, v))
+    queue: SimpleQueue[Dependency] = SimpleQueue()
+    for dependency in elements:
+        queue.put(dependency)
 
     last_name = None
     i = 0
     while True:
         try:
-            new, args = queue.get_nowait()
+            dependency = queue.get_nowait()
         except Empty:
             break
-        if args.issubset(available):
-            available.add(new)
-            order.append(new)
+        if dependency.required.issubset(available):
+            available.update(dependency.provided)
+            order.append(dependency.name)
+
         else:
-            if last_name == new:
-                order.append(new)
+            if last_name == dependency.name:
+                order.append(last_name)
                 break
-            queue.put((new, args))
-            last_name = new
+            queue.put(dependency)
+            last_name = dependency.name
         i += 1
 
         # Failure case
@@ -211,11 +226,13 @@ def _sort_dependencies(
             unsorted = []
             while True:
                 try:
-                    unsorted.append(queue.get_nowait()[0])
+                    unsorted.append(queue.get_nowait().name)
                 except Empty:
                     break
 
-            mod_to_args: dict[str, set[str]] = dict(elements)
+            mod_to_args: dict[str, set[str]] = {
+                dependency.name: dependency.required for dependency in elements
+            }
             missing = {k: mod_to_args[k].difference(available) for k in unsorted}
             raise CircularDependencyError(missing=missing)
     return order
@@ -303,7 +320,12 @@ class Model:
         to_sort = self._derived | self._reactions | self._surrogates
         order = _sort_dependencies(
             available=set(self._parameters) | set(self._variables) | {"time"},
-            elements=[(k, set(v.args)) for k, v in to_sort.items()],
+            elements=[
+                Dependency(name=k, required=set(v.args), provided={k})
+                if not isinstance(v, AbstractSurrogate)
+                else Dependency(name=k, required=set(v.args), provided=set(v.outputs))
+                for k, v in to_sort.items()
+            ],
         )
 
         # Split derived into parameters and variables
@@ -1227,6 +1249,7 @@ class Model:
         name: str,
         surrogate: AbstractSurrogate,
         args: list[str] | None = None,
+        outputs: list[str] | None = None,
         stoichiometries: dict[str, dict[str, float]] | None = None,
     ) -> Self:
         """Adds a surrogate model to the current instance.
@@ -1237,7 +1260,8 @@ class Model:
         Args:
             name (str): The name of the surrogate model.
             surrogate (AbstractSurrogate): The surrogate model instance to be added.
-            args: A list of arguments for the surrogate model.
+            args: Names of the values passed for the surrogate model.
+            outputs: Names of values produced by the surrogate model.
             stoichiometries: A dictionary mapping reaction names to stoichiometries.
 
         Returns:
@@ -1248,6 +1272,8 @@ class Model:
 
         if args is not None:
             surrogate.args = args
+        if outputs is not None:
+            surrogate.outputs = outputs
         if stoichiometries is not None:
             surrogate.stoichiometries = stoichiometries
 

mxlpy/npe/__init__.py

@@ -0,0 +1,38 @@
+"""Neural Process Estimation (NPE) module.
+
+This module provides classes and functions for estimating metabolic processes using
+neural networks. It includes functionality for both steady-state and time-course data.
+
+Classes:
+    TorchSteadyState: Class for steady-state neural network estimation.
+    TorchSteadyStateTrainer: Class for training steady-state neural networks.
+    TorchTimeCourse: Class for time-course neural network estimation.
+    TorchTimeCourseTrainer: Class for training time-course neural networks.
+
+Functions:
+    train_torch_steady_state: Train a PyTorch steady-state neural network.
+    train_torch_time_course: Train a PyTorch time-course neural network.
+"""
+
+from __future__ import annotations
+
+import contextlib
+
+with contextlib.suppress(ImportError):
+    from ._torch import (
+        TorchSteadyState,
+        TorchSteadyStateTrainer,
+        TorchTimeCourse,
+        TorchTimeCourseTrainer,
+        train_torch_steady_state,
+        train_torch_time_course,
+    )
+
+__all__ = [
+    "TorchSteadyState",
+    "TorchSteadyStateTrainer",
+    "TorchTimeCourse",
+    "TorchTimeCourseTrainer",
+    "train_torch_steady_state",
+    "train_torch_time_course",
+]