modelbase2 0.1.79__py3-none-any.whl → 0.2.0__py3-none-any.whl

modelbase2/experimental/codegen.py

@@ -0,0 +1,239 @@
+"""Module to export models as code."""
+
+import ast
+import inspect
+import warnings
+from collections.abc import Callable, Generator, Iterable, Iterator
+
+from modelbase2.model import Model
+from modelbase2.types import Derived
+
+__all__ = [
+    "DocstringRemover",
+    "IdentifierReplacer",
+    "ReturnRemover",
+    "conditional_join",
+    "generate_model_code_py",
+    "generate_modelbase_code",
+    "get_fn_source",
+    "handle_fn",
+]
+
+
+class IdentifierReplacer(ast.NodeTransformer):
+    """Replace identifiers in an AST."""
+
+    def __init__(self, mapping: dict[str, str]) -> None:
+        """Initialize the transformer with a mapping."""
+        self.mapping = mapping
+
+    def visit_Name(self, node: ast.Name) -> ast.Name:  # noqa: N802
+        """Replace the identifier with the mapped value."""
+        return ast.Name(
+            id=self.mapping.get(node.id, node.id),
+            ctx=node.ctx,
+        )
+
+
+class DocstringRemover(ast.NodeTransformer):
+    """Remove docstrings from an AST."""
+
+    def visit_Expr(self, node: ast.Expr) -> ast.Expr | None:  # noqa: N802
+        """Remove docstrings."""
+        if isinstance(const := node.value, ast.Constant) and isinstance(
+            const.value, str
+        ):
+            return None
+        return node
+
+
+class ReturnRemover(ast.NodeTransformer):
+    """Remove return statements from an AST."""
+
+    def visit_Return(self, node: ast.Return) -> ast.expr | None:  # noqa: N802
+        """Remove return statements."""
+        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)
+
+    argmap = dict(zip([i.arg for i in tree.args.args], args, strict=True))
+    tree = DocstringRemover().visit(tree)
+    tree = IdentifierReplacer(argmap).visit(tree)
+    tree = ReturnRemover().visit(tree)
+    return ast.unparse(tree.body)
+
+
+def conditional_join[T](
+    iterable: Iterable[T],
+    question: Callable[[T], bool],
+    true_pat: str,
+    false_pat: str,
+) -> str:
+    """Join an iterable, applying a pattern to each element based on a condition."""
+
+    def inner(it: Iterator[T]) -> Generator[str, None, None]:
+        yield str(next(it))
+        while True:
+            try:
+                el = next(it)
+                if question(el):
+                    yield f"{true_pat}{el}"
+                else:
+                    yield f"{false_pat}{el}"
+            except StopIteration:
+                break
+
+    return "".join(inner(iter(iterable)))
+
+
+def generate_modelbase_code(model: Model) -> str:
+    """Generate a modelbase model from a model."""
+    functions = {}
+
+    # Variables and parameters
+    variables = model.variables
+    parameters = model.parameters
+
+    # Derived
+    derived_source = []
+    for k, v in model.derived.items():
+        fn = v.fn
+        fn_name = fn.__name__
+        functions[fn_name] = inspect.getsource(fn)
+
+        derived_source.append(
+            f""".add_derived(
+                "{k}",
+                fn={fn_name},
+                args={v.args},
+            )"""
+        )
+
+    # Reactions
+    reactions_source = []
+    for k, v in model.reactions.items():
+        fn = v.fn
+        fn_name = fn.__name__
+        functions[fn_name] = inspect.getsource(fn)
+        reactions_source.append(
+            f"""        .add_reaction(
+                "{k}",
+                fn={fn_name},
+                args={v.args},
+                stoichiometry={v.stoichiometry},
+            )"""
+        )
+
+    # Surrogates
+    if len(model._surrogates) > 0:  # noqa: SLF001
+        warnings.warn(
+            "Generating code for Surrogates not yet supported.",
+            stacklevel=1,
+        )
+
+    # Combine all the sources
+    functions_source = "\n".join(functions.values())
+    source = [
+        "from modelbase2 import Model\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(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_model_code_py(model: Model) -> str:
+    """Transform the model into a single function, inlining the function calls."""
+    # Variables
+    variables = model.variables
+    variable_source = "    {} = y".format(", ".join(variables))
+
+    # Parameters
+    parameter_source = "\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)}")
+
+    # Reactions
+    reactions = []
+    for name, rxn in model.reactions.items():
+        reactions.append(f"    {name} = {handle_fn(rxn.fn, rxn.args)}")
+
+    # Stoichiometries
+    stoichiometries = {}
+    for rxn_name, rxn in model.reactions.items():
+        for cpd_name, factor in rxn.stoichiometry.items():
+            if isinstance(factor, Derived):
+                src = f"{handle_fn(factor.fn, factor.args)} * {rxn_name}"
+            elif factor == 1:
+                src = rxn_name
+            elif factor == -1:
+                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(
+            f"    d{variable}dt = {conditional_join(stoich, lambda x: x.startswith("-"), " ", " + ")}"
+        )
+
+    # Surrogates
+    if len(model._surrogates) > 0:  # noqa: SLF001
+        warnings.warn(
+            "Generating code for Surrogates not yet supported.",
+            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 "\n".join(source)

modelbase2/experimental/diff.py

@@ -0,0 +1,227 @@
+"""Diffing utilities for comparing models."""
+
+from dataclasses import dataclass, field
+from typing import Mapping
+
+from modelbase2.model import Model
+from modelbase2.types import Derived
+
+__all__ = ["DerivedDiff", "ModelDiff", "ReactionDiff", "model_diff", "soft_eq"]
+
+
+@dataclass
+class DerivedDiff:
+    """Difference between two derived variables."""
+
+    args1: list[str] = field(default_factory=list)
+    args2: list[str] = field(default_factory=list)
+
+
+@dataclass
+class ReactionDiff:
+    """Difference between two reactions."""
+
+    args1: list[str] = field(default_factory=list)
+    args2: list[str] = field(default_factory=list)
+    stoichiometry1: dict[str, float | Derived] = field(default_factory=dict)
+    stoichiometry2: dict[str, float | Derived] = field(default_factory=dict)
+
+
+@dataclass
+class ModelDiff:
+    """Difference between two models."""
+
+    missing_parameters: set[str] = field(default_factory=set)
+    missing_variables: set[str] = field(default_factory=set)
+    missing_reactions: set[str] = field(default_factory=set)
+    missing_surrogates: set[str] = field(default_factory=set)
+    missing_readouts: set[str] = field(default_factory=set)
+    missing_derived: set[str] = field(default_factory=set)
+    different_parameters: dict[str, tuple[float, float]] = field(default_factory=dict)
+    different_variables: dict[str, tuple[float, float]] = field(default_factory=dict)
+    different_reactions: dict[str, ReactionDiff] = field(default_factory=dict)
+    different_surrogates: dict[str, ReactionDiff] = field(default_factory=dict)
+    different_readouts: dict[str, DerivedDiff] = field(default_factory=dict)
+    different_derived: dict[str, DerivedDiff] = field(default_factory=dict)
+
+    def __str__(self) -> str:
+        """Return a human-readable string representation of the diff."""
+        content = ["Model Diff", "----------"]
+
+        # Parameters
+        if self.missing_parameters:
+            content.append(
+                "Missing Parameters: {}".format(", ".join(self.missing_parameters))
+            )
+        if self.different_parameters:
+            content.append("Different Parameters:")
+            for k, (v1, v2) in self.different_parameters.items():
+                content.append(f"  {k}: {v1} != {v2}")
+
+        # Variables
+        if self.missing_variables:
+            content.append(
+                "Missing Variables: {}".format(", ".join(self.missing_variables))
+            )
+        if self.different_variables:
+            content.append("Different Variables:")
+            for k, (v1, v2) in self.different_variables.items():
+                content.append(f"  {k}: {v1} != {v2}")
+
+        # Derived
+        if self.missing_derived:
+            content.append(
+                "Missing Derived: {}".format(", ".join(self.missing_derived))
+            )
+        if self.different_derived:
+            content.append("Different Derived:")
+            for k, diff in self.different_derived.items():
+                content.append(f"  {k}:")
+                if diff.args1 != diff.args2:
+                    content.append(f"    Args: {diff.args1} != {diff.args2}")
+
+        # Reactions
+        if self.missing_reactions:
+            content.append(
+                "Missing Reactions: {}".format(", ".join(self.missing_reactions))
+            )
+        if self.different_reactions:
+            content.append("Different Reactions:")
+            for k, diff in self.different_reactions.items():
+                content.append(f"  {k}:")
+                if diff.args1 != diff.args2:
+                    content.append(f"    Args: {diff.args1} != {diff.args2}")
+                if diff.stoichiometry1 != diff.stoichiometry2:
+                    content.append(
+                        f"    Stoichiometry: {diff.stoichiometry1} != {diff.stoichiometry2}"
+                    )
+
+        # Surrogates
+        if self.missing_surrogates:
+            content.append(
+                "Missing Surrogates: {}".format(", ".join(self.missing_surrogates))
+            )
+        if self.different_surrogates:
+            content.append("Different Surrogates:")
+            for k, diff in self.different_surrogates.items():
+                content.append(f"  {k}:")
+                if diff.args1 != diff.args2:
+                    content.append(f"    Args: {diff.args1} != {diff.args2}")
+                if diff.stoichiometry1 != diff.stoichiometry2:
+                    content.append(
+                        f"    Stoichiometry: {diff.stoichiometry1} != {diff.stoichiometry2}"
+                    )
+        return "\n".join(content)
+
+
+def _soft_eq_stoichiometries(
+    s1: Mapping[str, float | Derived], s2: Mapping[str, float | Derived]
+) -> bool:
+    """Check if two stoichiometries are equal, ignoring the functions."""
+    if s1.keys() != s2.keys():
+        return False
+
+    for k, v1 in s1.items():
+        v2 = s2[k]
+        if isinstance(v1, Derived):
+            if not isinstance(v2, Derived):
+                return False
+            if v1.args != v2.args:
+                return False
+        else:
+            if v1 != v2:
+                return False
+
+    return True
+
+
+def soft_eq(m1: Model, m2: Model) -> bool:
+    """Check if two models are equal, ignoring the functions."""
+    if m1._parameters != m2._parameters:  # noqa: SLF001
+        return False
+    if m1._variables != m2._variables:  # noqa: SLF001
+        return False
+    for k, d1 in m1._derived.items():  # noqa: SLF001
+        if (d2 := m2._derived.get(k)) is None:  # noqa: SLF001
+            return False
+        if d1.args != d2.args:
+            return False
+    for k, r1 in m1._readouts.items():  # noqa: SLF001
+        if (r2 := m2._readouts.get(k)) is None:  # noqa: SLF001
+            return False
+        if r1.args != r2.args:
+            return False
+    for k, v1 in m1._reactions.items():  # noqa: SLF001
+        if (v2 := m2._reactions.get(k)) is None:  # noqa: SLF001
+            return False
+        if v1.args != v2.args:
+            return False
+        if not _soft_eq_stoichiometries(v1.stoichiometry, v2.stoichiometry):
+            return False
+    for k, s1 in m1._surrogates.items():  # noqa: SLF001
+        if (s2 := m2._surrogates.get(k)) is None:  # noqa: SLF001
+            return False
+        if s1.args != s2.args:
+            return False
+        if s1.stoichiometries != s2.stoichiometries:
+            return False
+    return True
+
+
+def model_diff(m1: Model, m2: Model) -> ModelDiff:
+    """Compute the difference between two models."""
+    diff = ModelDiff()
+
+    for k, v1 in m1._parameters.items():  # noqa: SLF001
+        if (v2 := m2._parameters.get(k)) is None:  # noqa: SLF001
+            diff.missing_parameters.add(k)
+        elif v1 != v2:
+            diff.different_parameters[k] = (v1, v2)
+
+    for k, v1 in m1._variables.items():  # noqa: SLF001
+        if (v2 := m2._variables.get(k)) is None:  # noqa: SLF001
+            diff.missing_variables.add(k)
+        elif v1 != v2:
+            diff.different_variables[k] = (v1, v2)
+
+    for k, v1 in m1._readouts.items():  # noqa: SLF001
+        if (v2 := m2._readouts.get(k)) is None:  # noqa: SLF001
+            diff.missing_readouts.add(k)
+        elif v1.args != v2.args:
+            diff.different_readouts[k] = DerivedDiff(v1.args, v2.args)
+
+    for k, v1 in m1._derived.items():  # noqa: SLF001
+        if (v2 := m2._derived.get(k)) is None:  # noqa: SLF001
+            diff.missing_derived.add(k)
+        elif v1.args != v2.args:
+            diff.different_derived[k] = DerivedDiff(v1.args, v2.args)
+
+    for k, v1 in m1._reactions.items():  # noqa: SLF001
+        if (v2 := m2._reactions.get(k)) is None:  # noqa: SLF001
+            diff.missing_reactions.add(k)
+        else:
+            if v1.args != v2.args:
+                rxn_diff: ReactionDiff = diff.different_reactions.get(k, ReactionDiff())
+                rxn_diff.args1 = v1.args
+                rxn_diff.args2 = v2.args
+                diff.different_reactions[k] = rxn_diff
+            if v1.stoichiometry != v2.stoichiometry:
+                rxn_diff = diff.different_reactions.get(k, ReactionDiff())
+                rxn_diff.stoichiometry1 = dict(v1.stoichiometry)
+                rxn_diff.stoichiometry2 = dict(v2.stoichiometry)
+                diff.different_reactions[k] = rxn_diff
+
+    for k, v1 in m1._surrogates.items():  # noqa: SLF001
+        if (v2 := m2._surrogates.get(k)) is None:  # noqa: SLF001
+            diff.missing_surrogates.add(k)
+        else:
+            if v1.args != v2.args:
+                rxn_diff = diff.different_surrogates.get(k, ReactionDiff())
+                rxn_diff.args1 = v1.args
+                rxn_diff.args2 = v2.args
+            if v1.stoichiometries != v2.stoichiometries:
+                rxn_diff = diff.different_surrogates.get(k, ReactionDiff())
+                rxn_diff.stoichiometry1 = dict(v1.stoichiometries)  # type: ignore
+                rxn_diff.stoichiometry2 = dict(v2.stoichiometries)  # type: ignore
+
+    return diff

modelbase2/experimental/notes.md

@@ -0,0 +1,4 @@
+# Developer notes
+
+- Complete function expansion (e.g. inlining) other functions into the current function isn't possible in Python, as we can't expand statements
+- Same is true for assignment replacement, this only works for expressions