mxlpy 0.18.0__py3-none-any.whl → 0.20.0__py3-none-any.whl

mxlpy/integrators/int_scipy.py

@@ -2,13 +2,8 @@
 
 from __future__ import annotations
 
-from dataclasses import dataclass, field
-
-__all__ = [
-    "Scipy",
-]
-
 import copy
+from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, cast
 
 import numpy as np
@@ -20,6 +15,11 @@ if TYPE_CHECKING:
     from collections.abc import Callable
 
 
+__all__ = [
+    "Scipy",
+]
+
+
 @dataclass
 class Scipy:
     """Scipy integrator for solving ODEs.
@@ -108,10 +108,14 @@ class Scipy:
             rtol=self.rtol,
             method="LSODA",
         )
+
         if res.success:
-            self.t0 = time_points[-1]
-            self.y0 = res.y[:, -1]
-            return np.array(time_points, dtype=float), res.y.T
+            t = np.atleast_1d(np.array(res.t, dtype=float))
+            y = np.atleast_2d(np.array(res.y, dtype=float).T)
+
+            self.t0 = t[-1]
+            self.y0 = y[-1]
+            return t, y
         return None, None
 
     def integrate_to_steady_state(

mxlpy/label_map.py

@@ -24,11 +24,15 @@ import numpy as np
 
 from mxlpy.model import Model
 
-__all__ = ["LabelMapper"]
-
 if TYPE_CHECKING:
     from collections.abc import Callable, Mapping
 
+    from mxlpy.types import Derived
+
+__all__ = [
+    "LabelMapper",
+]
+
 
 def _total_concentration(*args: float) -> float:
     """Calculate sum of isotopomer concentrations.
@@ -552,7 +556,7 @@ class LabelMapper:
         for name, dp in self.model.derived_parameters.items():
             m.add_derived(name, fn=dp.fn, args=dp.args)
 
-        variables: dict[str, float] = {}
+        variables: dict[str, float | Derived] = {}
         for k, v in self.model.variables.items():
             if (isos := isotopomers.get(k)) is None:
                 variables[k] = v

mxlpy/linear_label_map.py

@@ -16,13 +16,15 @@ from typing import TYPE_CHECKING
 
 from mxlpy.model import Derived, Model
 
-__all__ = ["LinearLabelMapper"]
-
 if TYPE_CHECKING:
     from collections.abc import Mapping
 
     import pandas as pd
 
+__all__ = [
+    "LinearLabelMapper",
+]
+
 
 def _generate_isotope_labels(base_name: str, num_labels: int) -> list[str]:
     """Generate list of isotopomer names for a compound.

mxlpy/mc.py

@@ -42,6 +42,11 @@ from mxlpy.types import (
     TimeCourseByPars,
 )
 
+if TYPE_CHECKING:
+    from mxlpy.model import Model
+    from mxlpy.types import Array
+
+
 __all__ = [
     "ParameterScanWorker",
     "parameter_elasticities",
@@ -53,20 +58,6 @@ __all__ = [
     "variable_elasticities",
 ]
 
-if TYPE_CHECKING:
-    from mxlpy.model import Model
-    from mxlpy.types import Array
-
-__ALL__ = [
-    "steady_state",
-    "time_course",
-    "time_course_over_protocol",
-    "parameter_scan_ss",
-    "compound_elasticities",
-    "parameter_elasticities",
-    "response_coefficients",
-]
-
 
 class ParameterScanWorker(Protocol):
     """Protocol for the parameter scan worker function."""

mxlpy/mca.py

@@ -27,16 +27,16 @@ from mxlpy.parallel import parallelise
 from mxlpy.scan import _steady_state_worker
 from mxlpy.types import ResponseCoefficients
 
+if TYPE_CHECKING:
+    from mxlpy.model import Model
+    from mxlpy.types import IntegratorType
+
 __all__ = [
     "parameter_elasticities",
     "response_coefficients",
     "variable_elasticities",
 ]
 
-if TYPE_CHECKING:
-    from mxlpy.model import Model
-    from mxlpy.types import IntegratorType
-
 
 def _response_coefficient_worker(
     parameter: str,

mxlpy/meta/__init__.py

@@ -1,11 +1,13 @@
 """Metaprogramming facilities."""
 
+from __future__ import annotations
+
+from .codegen_latex import generate_latex_code
+from .codegen_modebase import generate_mxlpy_code
+from .codegen_py import generate_model_code_py
+
 __all__ = [
     "generate_latex_code",
     "generate_model_code_py",
     "generate_mxlpy_code",
 ]
-
-from .codegen_latex import generate_latex_code
-from .codegen_modebase import generate_mxlpy_code
-from .codegen_py import generate_model_code_py

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.variables,
+        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](