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

mxlpy/meta/codegen_latex.py

@@ -10,6 +10,11 @@ import sympy
 from mxlpy.meta.source_tools import fn_to_sympy
 from mxlpy.types import Derived, RateFn
 
+if TYPE_CHECKING:
+    from collections.abc import Callable, Mapping
+
+    from mxlpy import Model
+
 __all__ = [
     "TexExport",
     "TexReaction",
@@ -18,12 +23,6 @@ __all__ = [
     "get_model_tex_diff",
 ]
 
-if TYPE_CHECKING:
-    from collections.abc import Callable, Mapping
-
-    from mxlpy import Model
-
-
 cdot = r"\cdot"
 empty_set = r"\emptyset"
 left_right_arrows = r"\xleftrightharpoons{}"
@@ -94,8 +93,38 @@ def _escape_non_math(s: str) -> str:
     return s.replace("_", r"\_")
 
 
-def _fn_to_latex(fn: Callable, arg_names: list[str]) -> str:
-    return sympy.latex(fn_to_sympy(fn, _list_of_symbols(arg_names)))
+def _name_to_latex(s: str) -> str:
+    return _escape_non_math(_rename_latex(s))
+
+
+def _sympy_to_latex(expr: sympy.Expr) -> str:
+    return sympy.latex(
+        expr,
+        fold_frac_powers=True,
+        fold_func_brackets=True,
+        fold_short_frac=True,
+        mul_symbol="dot",
+    )
+
+
+def _fn_to_latex(
+    fn: Callable,
+    arg_names: list[str],
+    long_name_cutoff: int,
+) -> tuple[str, dict[str, str]]:
+    tex_names = [_mathrm(_name_to_latex(i)) for i in arg_names]
+    long_names = (
+        k
+        for k, k_orig in zip(tex_names, arg_names, strict=True)
+        if len(k_orig) >= long_name_cutoff
+    )
+    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_str = _sympy_to_latex(expr)
+    return fn_str, replacements
 
 
 def _table(
@@ -213,6 +242,12 @@ def _latex_list(rows: list[str]) -> str:
     return "\n\n".join(rows)
 
 
+def _latex_align(items: list[str]) -> str:
+    return rf"""\begin{{align*}}
+{"\n".join(items)}
+\end{{align*}}"""
+
+
 def _latex_list_as_sections(
     rows: list[tuple[str, str]], sec_fn: Callable[[str], str]
 ) -> str:
@@ -220,7 +255,7 @@ def _latex_list_as_sections(
         [
             "\n".join(
                 (
-                    sec_fn(_escape_non_math(name)),
+                    sec_fn(_name_to_latex(name)),
                     content,
                 )
             )
@@ -234,7 +269,7 @@ def _latex_list_as_bold(rows: list[tuple[str, str]]) -> str:
         [
             "\n".join(
                 (
-                    _bold(_escape_non_math(name)) + r"\\",
+                    _bold(_name_to_latex(name)) + r"\\",
                     content,
                     r"\vspace{20pt}",
                 )
@@ -244,33 +279,57 @@ def _latex_list_as_bold(rows: list[tuple[str, str]]) -> str:
     )
 
 
-def _stoichiometries_to_latex(stoich: Mapping[str, float | Derived]) -> str:
-    def optional_factor(k: str, v: float) -> str:
-        if v == 1:
-            return _mathrm(k)
-        if v == -1:
-            return f"-{_mathrm(k)}"
-        return f"{v} {cdot} {_mathrm(k)}"
+def _replacements_in_align(replacements: dict[str, str]) -> str:
+    reps = "\n".join(rf"&\qquad {v} :: {k} \\" for k, v in replacements.items())
 
-    def latex_for_empty(s: str) -> str:
-        if len(s) == 0:
-            return empty_set
-        return s
+    return rf"""&\quad \mathrm{{with}}\\
+{reps}\\"""
 
-    line = []
-    for k, v in stoich.items():
-        if isinstance(v, int | float):
-            line.append(optional_factor(k, v))
-        else:
-            line.append(_fn_to_latex(v.fn, [_rename_latex(i) for i in v.args]))
 
-    line_str = f"{line[0]}"
-    for element in line[1:]:
-        if element.startswith("-"):
-            line_str += f" {element}"
+def _diff_eq(name: str) -> str:
+    return rf"\frac{{d\left({name}\right)}}{{dt}}"
+
+
+def _optional_factor(k: str, v: float) -> str:
+    if v == 1:
+        return k
+    if v == -1:
+        return f"-{k}"
+    return f"{v} {cdot} {k}"
+
+
+def _stoichs_to_latex(
+    stoichs: Mapping[str, float | Derived],
+    long_name_cutoff: int,
+) -> tuple[str, dict[str, str]]:
+    replacements = {}
+    expr = sympy.Integer(0)
+
+    for rxn_name, rxn_stoich in stoichs.items():
+        rxn_name = _name_to_latex(rxn_name)  # noqa: PLW2901
+
+        if isinstance(rxn_stoich, Derived):
+            arg_names = [_mathrm(_name_to_latex(i)) for i in rxn_stoich.args]
+            long_names = (
+                k
+                for k, k_orig in zip(arg_names, rxn_stoich.args, strict=True)
+                if len(k_orig) >= long_name_cutoff
+            )
+            replacements.update(
+                {
+                    k: _name_to_latex(f"_x{i}")
+                    for i, k in enumerate(long_names, len(replacements))
+                }
+            )
+            sympy_fn = fn_to_sympy(
+                rxn_stoich.fn,
+                _list_of_symbols([replacements.get(k, k) for k in arg_names]),
+            )
+            expr = expr + sympy_fn * sympy.Symbol(rxn_name)  # type: ignore
         else:
-            line_str += f" + {element}"
-    return _math_il(line_str)
+            expr = expr + rxn_stoich * sympy.Symbol(rxn_name)  # type: ignore
+
+    return _sympy_to_latex(expr.subs(1.0, 1).simplify()), replacements
 
 
 @dataclass
@@ -333,7 +392,7 @@ class TexExport:
     variables: dict[str, float]
     derived: dict[str, Derived]
     reactions: dict[str, TexReaction]
-    stoichiometries: dict[str, Mapping[str, float | Derived]]
+    diff_eqs: dict[str, Mapping[str, float | Derived]]
 
     @staticmethod
     def _diff_parameters(
@@ -390,10 +449,8 @@ class TexExport:
             variables=self._diff_variables(self.variables, other.variables),
             derived=self._diff_derived(self.derived, other.derived),
             reactions=self._diff_reactions(self.reactions, other.reactions),
-            stoichiometries={
-                k: v
-                for k, v in other.stoichiometries.items()
-                if self.stoichiometries.get(k, {}) != v
+            diff_eqs={
+                k: v for k, v in other.diff_eqs.items() if self.diff_eqs.get(k, {}) != v
             },
         )
 
@@ -440,9 +497,9 @@ class TexExport:
                 )
                 for k, v in self.reactions.items()
             },
-            stoichiometries={
+            diff_eqs={
                 _add_gls_if_found(k): {gls.get(k2, k2): v2 for k2, v2 in v.items()}
-                for k, v in self.stoichiometries.items()
+                for k, v in self.diff_eqs.items()
             },
         )
 
@@ -466,7 +523,7 @@ class TexExport:
             headers=["Model name", "Initial concentration"],
             rows=[
                 [
-                    k,
+                    _name_to_latex(k),
                     f"{v:.2e}",
                 ]
                 for k, v in self.variables.items()
@@ -496,7 +553,7 @@ class TexExport:
         return _table(
             headers=["Parameter name", "Parameter value"],
             rows=[
-                [_math_il(_mathrm(_escape_non_math(_rename_latex(k)))), f"{v:.2e}"]
+                [_name_to_latex(k), f"{v:.2e}"]
                 for k, v in sorted(self.parameters.items())
             ],
             n_columns=2,
@@ -505,7 +562,10 @@ class TexExport:
             long_desc="Model parameters",
         )
 
-    def export_derived(self) -> str:
+    def export_derived(
+        self,
+        long_name_cutoff: int,
+    ) -> str:
         """Export derived quantities as LaTeX equations.
 
         Returns
@@ -523,16 +583,20 @@ class TexExport:
         True
 
         """
-        return _latex_list(
-            rows=[
-                _dmath(
-                    f"{_rename_latex(k)} = {_fn_to_latex(v.fn, [_rename_latex(i) for i in v.args])}"
-                )
-                for k, v in sorted(self.derived.items())
-            ]
-        )
+        rows = []
+        for k, v in sorted(self.derived.items()):
+            fn_str, repls = _fn_to_latex(
+                v.fn,
+                arg_names=v.args,
+                long_name_cutoff=long_name_cutoff,
+            )
+            rows.append(f"{_mathrm(_name_to_latex(k))} &= {fn_str} \\\\")
+            if repls:
+                rows.append(_replacements_in_align(repls))
 
-    def export_reactions(self) -> str:
+        return _latex_align(rows)
+
+    def export_reactions(self, long_name_cutoff: int) -> str:
         """Export reactions as LaTeX equations.
 
         Returns
@@ -550,16 +614,22 @@ class TexExport:
         True
 
         """
-        return _latex_list(
-            rows=[
-                _dmath(
-                    f"{_rename_latex(k)} = {_fn_to_latex(v.fn, [_rename_latex(i) for i in v.args])}"
-                )
-                for k, v in sorted(self.reactions.items())
-            ]
-        )
+        rows = []
+        for k, v in sorted(self.reactions.items()):
+            fn_str, repls = _fn_to_latex(
+                v.fn,
+                arg_names=v.args,
+                long_name_cutoff=long_name_cutoff,
+            )
+            rows.append(f"{_mathrm(_name_to_latex(k))} &= {fn_str} \\\\")
+            if repls:
+                rows.append(_replacements_in_align(repls))
+        return _latex_align(rows)
 
-    def export_stoichiometries(self) -> str:
+    def export_diff_eqs(
+        self,
+        long_name_cutoff: int,
+    ) -> str:
         """Export stoichiometries as LaTeX table.
 
         Returns
@@ -576,22 +646,20 @@ class TexExport:
         True
 
         """
-        return _table(
-            headers=["Rate name", "Stoichiometry"],
-            rows=[
-                [
-                    _escape_non_math(_rename_latex(k)),
-                    _stoichiometries_to_latex(v),
-                ]
-                for k, v in sorted(self.stoichiometries.items())
-            ],
-            n_columns=2,
-            label="model-stoichs",
-            short_desc="Model stoichiometries.",
-            long_desc="Model stoichiometries.",
-        )
+        rows = []
+        for var_name, stoich in sorted(self.diff_eqs.items()):
+            dxdt = _diff_eq(_mathrm(_name_to_latex(var_name)))
+            stoich_str, repls = _stoichs_to_latex(
+                stoich,
+                long_name_cutoff=long_name_cutoff,
+            )
+
+            rows.append(f"{dxdt} &= {stoich_str} \\\\")
+            if repls:
+                rows.append(_replacements_in_align(repls))
+        return _latex_align(rows)
 
-    def export_all(self) -> str:
+    def export_all(self, long_name_cutoff: int = 10) -> str:
         """Export all model parts as a complete LaTeX document section.
 
         Returns
@@ -626,28 +694,35 @@ class TexExport:
             sections.append(
                 (
                     "Derived",
-                    self.export_derived(),
+                    self.export_derived(
+                        long_name_cutoff=long_name_cutoff,
+                    ),
                 )
             )
         if len(self.reactions) > 0:
             sections.append(
                 (
                     "Reactions",
-                    self.export_reactions(),
+                    self.export_reactions(
+                        long_name_cutoff=long_name_cutoff,
+                    ),
                 )
             )
             sections.append(
                 (
-                    "Stoichiometries",
-                    self.export_stoichiometries(),
+                    "Differential equations",
+                    self.export_diff_eqs(
+                        long_name_cutoff=long_name_cutoff,
+                    ),
                 )
             )
-        return _latex_list_as_sections(sections, _subsubsection_)
+        return _latex_list_as_sections(sections, _subsection_)
 
     def export_document(
         self,
         author: str = "mxlpy",
         title: str = "Model construction",
+        long_name_cutoff: int = 10,
     ) -> str:
         r"""Export complete LaTeX document with all model components.
 
@@ -657,6 +732,8 @@ class TexExport:
             Name of the author for the document
         title
             Title for the document
+        long_name_cutoff
+            length of function argument names before they are shortened
 
         Returns
         -------
@@ -671,8 +748,8 @@ class TexExport:
         True
 
         """
-        content = self.export_all()
-        return rf"""\documentclass{{article}}
+        content = self.export_all(long_name_cutoff=long_name_cutoff)
+        return rf"""\documentclass[fleqn]{{article}}
 \usepackage[english]{{babel}}
 \usepackage[a4paper,top=2cm,bottom=2cm,left=2cm,right=2cm,marginparwidth=1.75cm]{{geometry}}
 \usepackage{{amsmath, amssymb, array, booktabs,
@@ -680,6 +757,7 @@ class TexExport:
             ragged2e, tabularx, titlesec, titling}}
 \newcommand{{\sectionbreak}}{{\clearpage}}
 \setlength{{\parindent}}{{0pt}}
+\allowdisplaybreaks
 
 \title{{{title}}}
 \date{{}} % clear date
@@ -692,18 +770,24 @@ class TexExport:
 
 
 def _to_tex_export(self: Model) -> TexExport:
+    diff_eqs = {}
+    for rxn_name, rxn in self.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()},
-        stoichiometries={k: v.stoichiometry for k, v in self.reactions.items()},
+        diff_eqs=diff_eqs,
     )
 
 
 def generate_latex_code(
     model: Model,
     gls: dict[str, str] | None = None,
+    long_name_cutoff: int = 10,
 ) -> str:
     """Export model as LaTeX document.
 
@@ -713,6 +797,8 @@ def generate_latex_code(
         The model to export
     gls
         Optional glossary mapping for renaming model components
+    long_name_cutoff
+        length of function argument names before they are shortened
 
     Returns
     -------
@@ -733,13 +819,18 @@ def generate_latex_code(
 
     """
     gls = default_init(gls)
-    return _to_tex_export(model).rename_with_glossary(gls).export_document()
+    return (
+        _to_tex_export(model)
+        .rename_with_glossary(gls)
+        .export_document(long_name_cutoff=long_name_cutoff)
+    )
 
 
 def get_model_tex_diff(
     m1: Model,
     m2: Model,
     gls: dict[str, str] | None = None,
+    long_name_cutoff: int = 10,
 ) -> str:
     """Create LaTeX diff showing changes between two models.
 
@@ -751,6 +842,8 @@ def get_model_tex_diff(
         Second model (compared against the base)
     gls
         Optional glossary mapping for renaming model components
+    long_name_cutoff
+        length of function argument names before they are shortened
 
     Returns
     -------
@@ -773,7 +866,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())}
+{((_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_modebase.py

@@ -13,7 +13,9 @@ from mxlpy.types import Derived
 if TYPE_CHECKING:
     from mxlpy.model import Model
 
-__all__ = ["generate_mxlpy_code"]
+__all__ = [
+    "generate_mxlpy_code",
+]
 
 
 def _list_of_symbols(args: list[str]) -> list[sympy.Symbol | sympy.Expr]:

mxlpy/meta/codegen_py.py

@@ -1,15 +1,23 @@
 """Module to export models as code."""
 
+from __future__ import annotations
+
 import warnings
-from collections.abc import Callable, Generator, Iterable, Iterator
+from typing import TYPE_CHECKING
 
 import sympy
 
 from mxlpy.meta.source_tools import fn_to_sympy, sympy_to_inline
-from mxlpy.model import Model
 from mxlpy.types import Derived
 
-__all__ = ["generate_model_code_py"]
+if TYPE_CHECKING:
+    from collections.abc import Callable, Generator, Iterable, Iterator
+
+    from mxlpy.model import Model
+
+__all__ = [
+    "generate_model_code_py",
+]
 
 
 def _conditional_join[T](

mxlpy/meta/source_tools.py

@@ -1,17 +1,21 @@
 """Tools for working with python source files."""
 
+from __future__ import annotations
+
 import ast
 import inspect
 import textwrap
-from collections.abc import Callable
 from dataclasses import dataclass
-from types import ModuleType
-from typing import Any, cast
+from typing import TYPE_CHECKING, Any, cast
 
 import dill
 import sympy
 from sympy.printing.pycode import pycode
 
+if TYPE_CHECKING:
+    from collections.abc import Callable
+    from types import ModuleType
+
 __all__ = [
     "Context",
     "fn_to_sympy",
@@ -35,7 +39,7 @@ class Context:
         symbols: dict[str, sympy.Symbol | sympy.Expr] | None = None,
         caller: Callable | None = None,
         parent_module: ModuleType | None = None,
-    ) -> "Context":
+    ) -> Context:
         """Update the context with new values."""
         return Context(
             symbols=self.symbols if symbols is None else symbols,

mxlpy/model.py

@@ -17,13 +17,13 @@ import numpy as np
 import pandas as pd
 
 from mxlpy import fns
-from mxlpy.types import (
-    AbstractSurrogate,
-    Array,
-    Derived,
-    Reaction,
-    Readout,
-)
+from mxlpy.types import AbstractSurrogate, Array, Derived, Reaction, Readout
+
+if TYPE_CHECKING:
+    from collections.abc import Iterable, Mapping
+    from inspect import FullArgSpec
+
+    from mxlpy.types import Callable, Param, RateFn, RetType
 
 __all__ = [
     "ArityMismatchError",
@@ -34,12 +34,6 @@ __all__ = [
     "ModelCache",
 ]
 
-if TYPE_CHECKING:
-    from collections.abc import Iterable, Mapping
-    from inspect import FullArgSpec
-
-    from mxlpy.types import Callable, Param, RateFn, RetType
-
 
 @dataclass
 class Dependency:
@@ -694,6 +688,21 @@ class Model:
 
         return self
 
+    def get_unused_parameters(self) -> set[str]:
+        """Get parameters which aren't used in the model."""
+        args = set()
+        for variable in self._variables.values():
+            if isinstance(variable, Derived):
+                args.update(variable.args)
+        for derived in self._derived.values():
+            args.update(derived.args)
+        for reaction in self._reactions.values():
+            args.update(reaction.args)
+        for surrogate in self._surrogates.values():
+            args.update(surrogate.args)
+
+        return set(self._parameters).difference(args)
+
     ##########################################################################
     # Variables
     ##########################################################################
@@ -1115,10 +1124,29 @@ class Model:
         args = self.get_dependent(variables=variables, time=time)
 
         stoich = copy.deepcopy(cache.stoich_by_cpds[variable])
-        for rxn, derived in cache.dyn_stoich_by_cpds[variable].items():
+        for rxn, derived in cache.dyn_stoich_by_cpds.get(variable, {}).items():
             stoich[rxn] = float(derived.fn(*(args[i] for i in derived.args)))
         return stoich
 
+    def get_raw_stoichiometries_of_variable(
+        self, variable: str
+    ) -> dict[str, float | Derived]:
+        """Retrieve the raw stoichiometry of a specific variable.
+
+        Examples:
+            >>> model.get_stoichiometries_of_variable("x1")
+                {"v1": -1, "v2": Derived(...)}
+
+        Args:
+            variable: The name of the variable for which to retrieve the stoichiometry.
+
+        """
+        stoichs: dict[str, dict[str, float | Derived]] = {}
+        for rxn_name, rxn in self._reactions.items():
+            for cpd_name, factor in rxn.stoichiometry.items():
+                stoichs.setdefault(cpd_name, {})[rxn_name] = factor
+        return stoichs[variable]
+
     @_invalidate_cache
     def add_reaction(
         self,

mxlpy/nn/_keras.py

@@ -1,12 +1,19 @@
-from typing import cast
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, cast
 
 import keras
 import pandas as pd
 from tqdm.keras import TqdmCallback
 
-from mxlpy.types import Array
+if TYPE_CHECKING:
+    from mxlpy.types import Array
 
-__all__ = ["LSTM", "MLP", "train"]
+__all__ = [
+    "LSTM",
+    "MLP",
+    "train",
+]
 
 
 def train(

mxlpy/nn/_torch.py

@@ -24,7 +24,13 @@ if TYPE_CHECKING:
 
     from mxlpy.types import Array
 
-__all__ = ["DefaultDevice", "LSTM", "LossFn", "MLP", "train"]
+__all__ = [
+    "DefaultDevice",
+    "LSTM",
+    "LossFn",
+    "MLP",
+    "train",
+]
 
 DefaultDevice = torch.device("cpu")