Merge pull request #221 from jtgreen/jtg/julia-mtk

Jtg/julia mtk

Author: finsberg

src/gotranx/cli/__init__.py

@@ -8,7 +8,7 @@
 from ..schemes import Scheme, get_scheme
 from ..codegen import PythonFormat, CFormat
 from ..codegen.base import Shape
-from . import gotran2c, gotran2py, gotran2julia, gotran2md
+from . import gotran2c, gotran2py, gotran2julia, gotran2md, gotran2mtk
 from . import utils
 
 app = typer.Typer()
@@ -515,6 +515,69 @@ def ode2julia(
     )
 
 
+@app.command()
+def ode2mtk(
+    fname: typing.Optional[Path] = typer.Argument(
+        None,
+        exists=True,
+        file_okay=True,
+        dir_okay=False,
+        writable=False,
+        readable=True,
+        resolve_path=True,
+    ),
+    outname: typing.Optional[str] = typer.Option(
+        None,
+        "-o",
+        "--outname",
+        help="Output name",
+    ),
+    remove_unused: bool = typer.Option(
+        False,
+        "--remove-unused",
+        help="Remove unused variables",
+    ),
+    version: bool = typer.Option(
+        None,
+        "--version",
+        callback=version_callback,
+        is_eager=True,
+        help="Show version",
+    ),
+    license: bool = typer.Option(
+        None,
+        "--license",
+        callback=license_callback,
+        is_eager=True,
+        help="Show license",
+    ),
+    config: typing.Optional[Path] = typer.Option(
+        None,
+        "-c",
+        "--config",
+        help="Read configuration options from a configuration file",
+    ),
+    verbose: bool = typer.Option(
+        False,
+        "--verbose",
+        "-v",
+        help="Verbose output",
+    ),
+):
+    if fname is None:
+        return typer.echo("No file specified")
+
+    config_data = utils.read_config(config)
+    verbose = config_data.get("verbose", verbose)
+    remove_unused = config_data.get("remove_unused", remove_unused)
+    gotran2mtk.main(
+        fname=fname,
+        outname=outname,
+        remove_unused=remove_unused,
+        verbose=verbose,
+    )
+
+
 @app.command()
 def list_schemes():
     from rich.console import Console

src/gotranx/cli/gotran2mtk.py

@@ -0,0 +1,41 @@
+from __future__ import annotations
+
+from pathlib import Path
+import logging
+import structlog
+
+from ..codegen.mtk import MTKCodeGenerator
+from ..load import load_ode
+from ..ode import ODE
+
+logger = structlog.get_logger()
+
+
+def get_code(
+    ode: ODE,
+    remove_unused: bool = False,
+) -> str:
+    """Generate ModelingToolkit.jl code for the ODE."""
+    codegen = MTKCodeGenerator(ode, remove_unused=remove_unused)
+    return codegen.generate()
+
+
+def main(
+    fname: Path,
+    outname: str | None = None,
+    remove_unused: bool = False,
+    verbose: bool = False,
+) -> None:
+    loglevel = logging.DEBUG if verbose else logging.INFO
+    structlog.configure(
+        wrapper_class=structlog.make_filtering_bound_logger(loglevel),
+    )
+    ode = load_ode(fname)
+    code = get_code(
+        ode,
+        remove_unused=remove_unused,
+    )
+    out = fname if outname is None else Path(outname)
+    out_name = out.with_suffix(suffix=".jl")
+    out_name.write_text(code)
+    logger.info(f"Wrote {out_name}")

src/gotranx/codegen/__init__.py

@@ -9,6 +9,7 @@
 from .base import CodeGenerator, Func, RHSArgument, SchemeArgument
 from .ode import GotranODECodePrinter, BaseGotranODECodePrinter
 from .julia import JuliaCodeGenerator, GotranJuliaCodePrinter
+from .mtk import MTKCodeGenerator
 
 __all__ = [
     "base",
@@ -30,6 +31,5 @@
     "JuliaCodeGenerator",
     "GotranJuliaCodePrinter",
     "JaxCodeGenerator",
-    "JuliaCodeGenerator",
-    "GotranJuliaCodePrinter",
+    "MTKCodeGenerator",
 ]

src/gotranx/codegen/mtk.py

@@ -0,0 +1,103 @@
+from __future__ import annotations
+
+from typing import Iterable
+
+import sympy
+import structlog
+
+from .base import CodeGenerator, Func, RHSArgument, SchemeArgument
+from .julia import GotranJuliaCodePrinter
+from .. import templates
+from ..ode import ODE
+from .. import atoms
+
+logger = structlog.get_logger()
+
+
+class MTKCodeGenerator(CodeGenerator):
+    def __init__(self, ode: ODE, remove_unused: bool = False) -> None:
+        super().__init__(ode, remove_unused=remove_unused)
+        self._printer = GotranJuliaCodePrinter()
+
+    @property
+    def printer(self):
+        return self._printer
+
+    @property
+    def template(self):
+        return templates.mtk
+
+    # The following methods satisfy the abstract interface but are not used
+    def _rhs_arguments(self, order: RHSArgument | str = RHSArgument.tsp, const_states: bool = True):
+        dummy = sympy.IndexedBase("dummy")
+        return Func(
+            arguments=[],
+            states=dummy,
+            parameters=dummy,
+            values=dummy,
+            values_type="",
+        )
+
+    def _scheme_arguments(
+        self,
+        order: SchemeArgument | str = SchemeArgument.stdp,
+        const_states: bool = True,
+    ):
+        dummy = sympy.IndexedBase("dummy")
+        return Func(
+            arguments=[],
+            states=dummy,
+            parameters=dummy,
+            values=dummy,
+            values_type="",
+        )
+
+    def _format_expr(self, expr) -> str:
+        return self.printer.doprint(expr)
+
+    def _observed(self, assignments: Iterable[atoms.Assignment]) -> list[str]:
+        observed = []
+        for assignment in assignments:
+            if isinstance(assignment, atoms.Intermediate):
+                observed.append(f"{assignment.name} ~ {self._format_expr(assignment.expr)}")
+        return observed
+
+    def _equations(self, derivatives: Iterable[atoms.StateDerivative]) -> list[str]:
+        equations = []
+        for derivative in derivatives:
+            lhs = f"D({derivative.state.name})"
+            rhs = self._format_expr(derivative.expr)
+            equations.append(f"{lhs} ~ {rhs}")
+        return equations
+
+    def generate(self, remove_unused: bool | None = None) -> str:
+        """Generate ModelingToolkit.jl code."""
+        if remove_unused is None:
+            remove_unused = self.remove_unused
+
+        param_names = [p.name for p in self.ode.parameters if self._condition(p.name)]
+        missing_names = [n for n in self.ode.missing_variables.keys() if self._condition(n)]
+        extra_params = [n for n in missing_names if n not in param_names]
+
+        state_names = [s.name for s in self.ode.sorted_states() if self._condition(s.name)]
+
+        assignments = self.ode.sorted_assignments(remove_unused=remove_unused)
+        observed = self._observed(assignments)
+        derivatives = [a for a in assignments if isinstance(a, atoms.StateDerivative)]
+        equations = self._equations(derivatives)
+
+        state_defaults = [(s.name, self._format_expr(s.value)) for s in self.ode.sorted_states()]
+        param_defaults = [(p.name, self._format_expr(p.value)) for p in self.ode.parameters]
+        param_defaults.extend((n, "0.0") for n in extra_params)
+
+        parts = [
+            self.template.header(),
+            self.template.parameters_block(param_names, extra_params),
+            self.template.states_block(state_names),
+            self.template.observed_block(observed),
+            self.template.equations_block(equations),
+            self.template.defaults_block(state_defaults, param_defaults),
+            self.template.ode_system(self.ode.name, state_names, param_names + extra_params),
+        ]
+
+        return "\n".join(parts)

src/gotranx/templates/__init__.py

@@ -5,6 +5,7 @@
 from . import python
 from . import jax
 from . import julia
+from . import mtk
 from . import markdown
 
 
@@ -195,4 +196,4 @@ def method(
         """
 
 
-__all__ = ["c", "python", "jax", "julia", "markdown", "Template"]
+__all__ = ["c", "python", "jax", "julia", "markdown", "mtk", "Template"]

src/gotranx/templates/mtk.py

@@ -0,0 +1,92 @@
+from __future__ import annotations
+
+from textwrap import dedent, indent
+
+
+def header() -> str:
+    return dedent(
+        """
+        using ModelingToolkit
+        """,
+    )
+
+
+def parameters_block(params: list[str], extra_params: list[str]) -> str:
+    names = " ".join(["t"] + params + extra_params)
+    return dedent(
+        f"""
+        @parameters {names}
+        D = Differential(t)
+        """,
+    )
+
+
+def states_block(states: list[str]) -> str:
+    names = " ".join([f"{s}(t)" for s in states])
+    return dedent(
+        f"""
+        @variables {names}
+        """,
+    )
+
+
+def observed_block(entries: list[str]) -> str:
+    if not entries:
+        return "observed = []"
+
+    body = indent(",\n".join(entries), "    ")
+    return dedent(
+        f"""
+        observed = [
+        {body}
+        ]
+        """,
+    )
+
+
+def equations_block(entries: list[str]) -> str:
+    body = indent(",\n".join(entries), "    ")
+    return dedent(
+        f"""
+        eqs = [
+        {body}
+        ]
+        """,
+    )
+
+
+def defaults_block(
+    state_defaults: list[tuple[str, str]], param_defaults: list[tuple[str, str]]
+) -> str:
+    pairs = []
+    for name, value in state_defaults + param_defaults:
+        pairs.append(f"{name} => {value}")
+
+    if not pairs:
+        return "defaults = Dict()"
+
+    body = indent(",\n".join(pairs), "    ")
+    return dedent(
+        f"""
+        defaults = Dict(
+        {body}
+        )
+        """,
+    )
+
+
+def ode_system(name: str, state_names: list[str], param_names: list[str]) -> str:
+    states = ", ".join(state_names)
+    params = ", ".join(param_names)
+    return dedent(
+        f"""
+        @named {name} = ODESystem(
+            eqs,
+            t,
+            [{states}],
+            [{params}],
+            observed = observed,
+            defaults = defaults,
+        )
+        """,
+    )