sindy-exp 0.2.0__py3-none-any.whl

sindy_exp/_dysts_to_sympy.py

@@ -0,0 +1,452 @@
+"""Utilities to convert a `dysts` dynamical system object's rhs to SymPy.
+
+This module inspects the source of an object's RHS method (by default
+named ``rhs``), parses the function using ``ast``, and converts the
+returned expression(s) into SymPy expressions.
+
+The conversion is intentionally conservative and aims to handle common
+patterns used in simple rhs implementations, e.g. returning a tuple/list
+of arithmetic expressions, using indexing into a state vector (``x[0]``),
+and calls to common ``numpy``/``math`` functions (``np.sin``, ``math.exp``, ...).
+
+Limitations:
+- It does not execute arbitrary code from the inspected function.
+- Complex control flow, loops, or non-trivial Python constructs may not
+  be fully supported.
+
+Example
+-------
+from dysts.flows import Lorenz
+from inspect_to_sympy import object_to_sympy_rhs
+
+lor = Lorenz()
+symbols, exprs, lambda_rhs = object_to_sympy_rhs(lor)
+# `symbols` is a list of SymPy symbols for the state vector
+# `exprs` is a list of SymPy expressions for the RHS
+# `lambda_rhs` is a SymPy Lambda mapping state symbols -> rhs expressions
+"""
+
+from __future__ import annotations
+
+import ast
+import inspect
+import textwrap
+from typing import Any, Callable, Dict, List, Tuple
+
+import numpy as np
+import sympy as sp
+from dysts.base import BaseDyn
+
+
+def _is_name(node: ast.AST, name: str) -> bool:
+    return isinstance(node, ast.Name) and node.id == name
+
+
+class _ASTToSympy(ast.NodeVisitor):
+    def __init__(
+        self,
+        state_name: str,
+        state_symbols: List[sp.Symbol],
+        locals_map: Dict[str, Any],
+    ):
+        self.state_name = state_name
+        self.state_symbols = state_symbols
+        self.locals = dict(locals_map)
+
+    def generic_visit(self, node):
+        raise NotImplementedError(f"AST node not supported: {node!r}")
+
+    def visit_Constant(self, node: ast.Constant):
+        return sp.sympify(node.value)
+
+    def visit_Num(self, node: ast.Num):
+        return sp.sympify(node.n)
+
+    def visit_Name(self, node: ast.Name):
+        if node.id in self.locals:
+            return self.locals[node.id]
+        return sp.Symbol(node.id)
+
+    def visit_Tuple(self, node: ast.Tuple):
+        elems = []
+        for elt in node.elts:
+            val = self.visit(elt)
+            if isinstance(val, (list, tuple)):
+                elems.extend(list(val))
+            else:
+                elems.append(val)
+        return tuple(elems)
+
+    def visit_List(self, node: ast.List):
+        elems = []
+        for elt in node.elts:
+            val = self.visit(elt)
+            if isinstance(val, (list, tuple)):
+                elems.extend(list(val))
+            else:
+                elems.append(val)
+        return elems
+
+    def visit_Starred(self, node: ast.Starred):
+        # Handle starred expressions like `*x` in list/tuple literals.
+        # If the starred value is the state vector name, expand to state symbols.
+        if isinstance(node.value, ast.Name) and node.value.id == self.state_name:
+            return tuple(self.state_symbols)
+        # Otherwise, evaluate the value and if it is a sequence, return its items
+        val = self.visit(node.value)
+        if isinstance(val, (list, tuple)):
+            return tuple(val)
+        raise NotImplementedError(
+            "Unsupported starred expression; cannot expand non-iterable"
+        )
+
+    def visit_BinOp(self, node: ast.BinOp):
+        left = self.visit(node.left)
+        right = self.visit(node.right)
+        if isinstance(node.op, ast.Add):
+            return left + right
+        if isinstance(node.op, ast.Sub):
+            return left - right
+        if isinstance(node.op, ast.Mult):
+            return left * right
+        if isinstance(node.op, ast.Div):
+            return left / right
+        if isinstance(node.op, ast.Pow):
+            return left**right
+        if isinstance(node.op, ast.Mod):
+            return left % right
+        raise NotImplementedError(f"Binary op not supported: {node.op!r}")
+
+    def visit_UnaryOp(self, node: ast.UnaryOp):
+        operand = self.visit(node.operand)
+        if isinstance(node.op, ast.USub):
+            return -operand
+        if isinstance(node.op, ast.UAdd):
+            return +operand
+        raise NotImplementedError(f"Unary op not supported: {node.op!r}")
+
+    def visit_Call(self, node: ast.Call):
+        # Determine function name
+        func = node.func
+        func_name = None
+        mod_name = None
+
+        if isinstance(func, ast.Name):
+            func_name = func.id
+        elif isinstance(func, ast.Attribute):
+            # e.g. np.sin or math.exp
+            if isinstance(func.value, ast.Name):
+                mod_name = func.value.id
+            func_name = func.attr
+        else:
+            raise NotImplementedError(
+                f"Call to unsupported func node: {ast.dump(func)}"
+            )
+
+        # Map common numpy/math functions to sympy
+        func_map = {
+            "sin": sp.sin,
+            "cos": sp.cos,
+            "tan": sp.tan,
+            "exp": sp.exp,
+            "log": sp.log,
+            "sqrt": sp.sqrt,
+            "abs": sp.Abs,
+            "atan": sp.atan,
+            "asin": sp.asin,
+            "acos": sp.acos,
+        }
+
+        args = [self.visit(a) for a in node.args]
+
+        # Special-case array constructors: return underlying list/tuple
+        if func_name in ("array", "asarray") and mod_name in ("np", "numpy"):
+            # expect a single positional arg that's a list/tuple
+            if len(args) == 1:
+                return args[0]
+
+        if func_name in func_map:
+            return func_map[func_name](*args)
+
+        # Unknown function: create a Sympy Function
+        symf = sp.Function(func_name)
+        return symf(*args)
+
+    def visit_Subscript(self, node: ast.Subscript):
+        # Support patterns like x[0] where x is the state vector name
+        value = node.value
+        # handle simple constant index
+        if _is_name(value, self.state_name):
+            # Python >=3.9: slice is directly the node.slice
+            idx_node = node.slice
+            if isinstance(idx_node, ast.Constant):
+                idx = idx_node.value
+            else:
+                raise NotImplementedError(
+                    "Only constant indices into state vector supported"
+                )
+            return self.state_symbols[idx]
+
+        # If it's something else, try to evaluate generically
+        base = self.visit(value)
+        # slice may be constant
+        if isinstance(node.slice, ast.Constant):
+            key = node.slice.value
+            return base[key]
+        raise NotImplementedError("Unsupported subscript pattern")
+
+
+def _numeric_consistency_check(
+    dysts_flow: BaseDyn,
+    rhsfunc: Callable,
+    arg_names: List[str],
+    state_names: List[str],
+    vector_mode: bool,
+    sys_dim: int,
+    lambda_rhs: sp.Lambda,
+) -> None:
+    """Compare the original dysts rhs function to the SymPy-derived lambda.
+
+    Raises a RuntimeError if they disagree.
+    """
+    # default to nonnegative support (e.g. Lotka volterra)
+    random_state = np.random.standard_exponential(size=sys_dim)
+
+    # Construct call arguments for the original function (bound method).
+    call_args = []
+    for name in arg_names:
+        if name == "self":
+            continue
+        if name in state_names and not vector_mode:
+            idx = state_names.index(name)
+            call_args.append(random_state[idx])
+        elif name in state_names and vector_mode:
+            call_args.append(np.asarray(random_state, dtype=float))
+        elif name == "t":
+            call_args.append(float(np.random.standard_normal(size=())))
+        else:
+            call_args.append(dysts_flow.params[name])
+
+    dysts_val = rhsfunc(*call_args)
+    orig_arr = np.asarray(dysts_val, dtype=float).ravel()
+
+    sym_val = lambda_rhs(*tuple(random_state))
+    sym_arr = np.asarray(sym_val, dtype=float).ravel()
+
+    if orig_arr.shape != sym_arr.shape:
+        raise RuntimeError(
+            f"_rhs shape {orig_arr.shape} != sympy shape {sym_arr.shape}"
+        )
+
+    if not np.allclose(orig_arr, sym_arr, rtol=1e-6, atol=1e-9):
+        raise RuntimeError("Numeric mismatch between original and sympy conversion.")
+
+
+def dynsys_to_sympy(
+    obj: Any, func_name: str = "_rhs"
+) -> Tuple[List[sp.Symbol], List[sp.Expr], sp.Lambda]:
+    """Inspect ``obj`` for a method named ``func_name`` and return a SymPy
+    representation of its RHS.
+
+    Returns:
+        a tuple ``(state_symbols, exprs, lambda_rhs)`` where ``state_symbols``
+    is a list of SymPy symbols for the state vector, ``exprs`` is a list of
+    SymPy expressions for the RHS components, and ``lambda_rhs`` is a SymPy
+    Lambda mapping the state symbols to the RHS vector.
+
+    Example:
+
+    >>> from dysts.flows import Lorenz
+    >>> from inspect_to_sympy import dynsys_to_sympy
+    >>> lor = Lorenz()
+    >>> symbols, exprs, lambda_rhs = dynsys_to_sympy(lor)
+    >>> print(lor._rhs(1, 2, 3, t=0.0, **lor.params))
+    (10, 23, -6.0009999999999994)
+
+    >>> print(tuple(lambda_rhs(1, 2, 3)))
+    (10, 23, -6.00100000000000)
+
+    """
+
+    if not hasattr(obj, func_name):
+        raise AttributeError(f"Object has no attribute {func_name!r}")
+
+    func = getattr(obj, func_name)
+    src = inspect.getsource(func)
+    src = textwrap.dedent(src)
+
+    parsed = ast.parse(src)
+
+    # Find first FunctionDef
+    fndef = None
+    for node in parsed.body:
+        if isinstance(node, ast.FunctionDef):
+            fndef = node
+            break
+    if fndef is None:
+        raise RuntimeError("No function definition found in source")
+
+    # Determine state argument names. Common dysts signature:
+    # (self, *states, t, *parameters). Prefer obj.dimension when available.
+    arg_names = [a.arg for a in fndef.args.args]
+    if len(arg_names) == 0:
+        raise RuntimeError("Function has no arguments")
+
+    start_idx = 0
+    if arg_names[0] == "self":
+        start_idx = 1
+
+    vector_mode = False
+    state_args: List[str]
+    t_idx = None
+    if "t" in arg_names:
+        t_idx = arg_names.index("t")
+
+    if hasattr(obj, "dimension") and isinstance(getattr(obj, "dimension"), int):
+        n_state = int(getattr(obj, "dimension"))
+        if t_idx is not None:
+            potential = arg_names[start_idx:t_idx]
+            if len(potential) >= n_state:
+                state_args = potential[:n_state]
+            else:
+                state_args = [arg_names[start_idx]]
+                vector_mode = True
+        else:
+            potential = arg_names[start_idx:]
+            if len(potential) >= n_state:
+                state_args = potential[:n_state]
+            else:
+                state_args = [arg_names[start_idx]]
+                vector_mode = True
+    else:
+        if t_idx is not None:
+            state_args = arg_names[start_idx:t_idx]
+            if len(state_args) == 0:
+                state_args = [arg_names[start_idx]]
+                vector_mode = True
+            elif len(state_args) == 1:
+                # single name could be vector or scalar; assume vector-mode
+                vector_mode = True
+        else:
+            state_args = [arg_names[start_idx]]
+            vector_mode = True
+
+    # If vector_mode, inspect AST for subscript/index usage or tuple unpacking
+    if vector_mode:
+        state_name = state_args[0]
+        max_index = -1
+        unpack_size = None
+        for node in ast.walk(fndef):
+            if (
+                isinstance(node, ast.Subscript)
+                and isinstance(node.value, ast.Name)
+                and node.value.id == state_name
+            ):
+                sl = node.slice
+                if isinstance(sl, ast.Constant) and isinstance(sl.value, int):
+                    if sl.value > max_index:
+                        max_index = sl.value
+            if isinstance(node, ast.Assign):
+                if isinstance(node.value, ast.Name) and node.value.id == state_name:
+                    targets = node.targets
+                    if len(targets) == 1 and isinstance(
+                        targets[0], (ast.Tuple, ast.List)
+                    ):
+                        unpack_size = len(targets[0].elts)
+
+        if unpack_size is not None:
+            n_state = unpack_size
+        elif max_index >= 0:
+            n_state = max_index + 1
+        else:
+            n_state = int(getattr(obj, "dimension", 3))
+
+        state_symbols = [sp.Symbol(f"x{i}") for i in range(n_state)]
+        primary_state_name = state_name
+    else:
+        # individual state args -> use their arg names as symbol names
+        state_symbols = [sp.Symbol(n) for n in state_args]
+        primary_state_name = state_args[0] if len(state_args) > 0 else "x"
+
+    # Build locals mapping from known state arg names and parameters
+    locals_map: Dict[str, Any] = {}
+    for i, name in enumerate(state_args):
+        if i < len(state_symbols):
+            locals_map[name] = state_symbols[i]
+
+    # map parameters (if present) to numeric values or symbols
+    if hasattr(obj, "parameters") and isinstance(getattr(obj, "parameters"), dict):
+        params = getattr(obj, "parameters")
+        if t_idx is not None:
+            param_arg_names = arg_names[t_idx + 1 :]
+        else:
+            param_arg_names = []
+        for pname in param_arg_names:
+            if pname in params:
+                locals_map[pname] = sp.sympify(params[pname])
+            else:
+                locals_map[pname] = sp.Symbol(pname)
+
+    converter = _ASTToSympy(primary_state_name, state_symbols, locals_map)
+
+    return_expr = None
+    # Walk through function body statements, handle Assign and Return
+    for stmt in fndef.body:
+        if isinstance(stmt, ast.Assign):
+            # only simple single-target assignments supported
+            if len(stmt.targets) != 1:
+                raise ValueError("Only single-target assignments supported")
+            target = stmt.targets[0]
+            if isinstance(target, ast.Name):
+                value_expr = converter.visit(stmt.value)
+                locals_map[target.id] = value_expr
+            elif (
+                isinstance(target, (ast.Tuple, ast.List))
+                and isinstance(stmt.value, ast.Name)
+                and stmt.value.id == state_name
+            ):
+                # unpacking like a,b,c = x -> map names to state symbols
+                for i, elt in enumerate(target.elts):
+                    if isinstance(elt, ast.Name):
+                        locals_map[elt.id] = state_symbols[i]
+        elif isinstance(stmt, ast.Return):
+            return_expr = stmt.value
+
+    if return_expr is None:
+        # maybe last statement is an Expr with list construction;
+        # try to find a Return node deep
+        for node in ast.walk(fndef):
+            if isinstance(node, ast.Return):
+                return_expr = node.value
+                break
+
+    if return_expr is None:
+        raise RuntimeError("No return expression found in function body")
+
+    # Refresh converter with updated locals
+    converter = _ASTToSympy(primary_state_name, state_symbols, locals_map)
+    rhs_val = converter.visit(return_expr)
+
+    if isinstance(rhs_val, (list, tuple)):
+        exprs = list(rhs_val)
+    else:
+        # single expression: treat as 1-dim RHS
+        exprs = [rhs_val]
+
+    lambda_rhs = sp.Lambda(tuple(state_symbols), sp.Matrix(exprs))
+
+    # Run numeric consistency guard (raises on mismatch)
+    _numeric_consistency_check(
+        obj,
+        func,
+        arg_names,
+        state_args,
+        vector_mode,
+        len(state_symbols),
+        lambda_rhs,
+    )
+
+    return state_symbols, exprs, lambda_rhs
+
+
+__all__ = ["dynsys_to_sympy"]