demathpy 0.0.2__py3-none-any.whl → 0.1.1__py3-none-any.whl

demathpy/__init__.py

@@ -2,32 +2,22 @@
 
 from .symbols import normalize_symbols, normalize_lhs
 from .pde import (
+    PDE,
     normalize_pde,
     init_grid,
-    sample_gradient,
     parse_pde,
-    evaluate_rhs,
-    evaluate_scalar,
     step_pdes,
-    evaluate_rhs_compiled,
-    evaluate_scalar_compiled,
-    step_compiled_pdes,
 )
-from .ode import robust_parse, parse_odes_to_function
+from .ode import ODE, parse_ode
 
 __all__ = [
+    "PDE",
+    "ODE",
     "normalize_symbols",
     "normalize_lhs",
     "normalize_pde",
     "init_grid",
-    "sample_gradient",
     "parse_pde",
-    "evaluate_rhs",
-    "evaluate_scalar",
     "step_pdes",
-    "evaluate_rhs_compiled",
-    "evaluate_scalar_compiled",
-    "step_compiled_pdes",
-    "robust_parse",
-    "parse_odes_to_function",
+    "parse_ode",
 ]

demathpy/ode.py

@@ -1,131 +1,377 @@
-import json
+"""
+Ordinary Differential Equation (ODE) utilities.
+
+Provides a class-based ODE solver similar to the PDE module, 
+but for systems dependent only on time (t).
+"""
+
+from __future__ import annotations
+
+from typing import Dict, List, Tuple, Any, Optional
+
 import re
-import sympy
+import json
 import numpy as np
-from sympy.parsing.sympy_parser import parse_expr, standard_transformations, implicit_multiplication_application, convert_xor
 
+from .symbols import normalize_symbols, normalize_lhs
 
-def _convert_ternary(expr: str) -> str:
-    """
-    Convert a single C-style ternary (cond ? a : b) into SymPy Piecewise.
-    Supports one level (no nesting).
-    """
-    if "?" not in expr or ":" not in expr:
-        return expr
-
-    # naive split for single ternary
-    # pattern: <cond> ? <a> : <b>
-    parts = expr.split("?")
-    if len(parts) != 2:
-        return expr
-    cond = parts[0].strip()
-    rest = parts[1]
-    if ":" not in rest:
-        return expr
-    a, b = rest.split(":", 1)
-    a = a.strip()
-    b = b.strip()
-    return f"Piecewise(({a}, {cond}), ({b}, True))"
-
-
-def robust_parse(expr_str):
+
+def _preprocess_expr(expr: str) -> str:
+    expr = (expr or "").strip()
+    expr = re.sub(r"\(\s*approx\s*\)", "", expr, flags=re.IGNORECASE)
+    expr = re.sub(r"\bapprox\b", "", expr, flags=re.IGNORECASE)
+    if "=" in expr:
+        expr = expr.split("=")[-1]
+    expr = normalize_symbols(expr)
+    return expr.strip()
+
+
+def normalize_ode(ode: str) -> str:
+    return (ode or "").strip()
+
+
+def parse_ode(ode: str) -> Tuple[str, int, str, str]:
     """
-    Parses a string into a SymPy expression with relaxed syntax rules:
-    - Implicit multiplication (5x -> 5*x)
-    - Caret for power (x^2 -> x**2)
-    - Aliases 'y' to 'z' for 2D convenience
+    Returns (var, order, lhs_coeff_expr, rhs_expr).
+    Parses "dy/dt = -y" or "d^2y/dt^2 = -y".
     """
-    if not isinstance(expr_str, str):
-        return sympy.sympify(expr_str) 
+    ode = normalize_ode(ode)
+    if "=" not in ode:
+        # Assume implicit "dy/dt =" ?? No, safer to default to u, 1st order
+        return "u", 1, "1", ode
+
+    lhs, rhs = ode.split("=", 1)
+    lhs = normalize_lhs(lhs.strip())
+    rhs = rhs.strip()
+
+    def _extract_coeff(lhs_expr: str, deriv_expr: str) -> str:
+        coeff = lhs_expr.replace(deriv_expr, "").strip()
+        if not coeff:
+            return "1"
+        coeff = coeff.strip("*")
+        return _preprocess_expr(coeff) or "1"
+    
+    # Updated regex for ODE derivatives: dy/dt, d^2x/dt^2
+    pattern = r"(?:∂|d)\s*(?:\^?(\d+))?\s*([a-zA-Z_]\w*)\s*/\s*(?:∂|d)t\s*(?:\^?(\d+))?"
+    m = re.search(pattern, lhs)
+    
+    if m:
+        ord1 = m.group(1)
+        var = m.group(2)
+        ord2 = m.group(3)
         
-    transformations = (standard_transformations + (implicit_multiplication_application, convert_xor))
+        order = 1
+        if ord1:
+            order = int(ord1)
+        elif ord2:
+            order = int(ord2)
+            
+        coeff = _extract_coeff(lhs, m.group(0))
+        return var, order, coeff, _preprocess_expr(rhs)
+
+    # Fallback if no dy/dt found on lhs?
+    # Maybe user wrote "y' = ..." (not supported by regex yet)
+    return "u", 1, "1", _preprocess_expr(rhs)
+
+
+class ODE:
+    equation: str
+    desc: str
+
+    u: np.ndarray 
+    u_shape: List[str] # ["x"] or ["x", "y"] for vector systems
     
-    # Define symbols and alias y -> z
-    x, z, vx, vz, t, pid = sympy.symbols('x z vx vz t id')
-    local_dict = {
-        'x': x, 'z': z, 'y': z, 'vx': vx, 'vz': vz, 't': t, 'id': pid,
-        'pi': sympy.pi, 'e': sympy.E
-    }
-
-    # Ensure common functions are recognized (Abs not abs)
-    local_dict.update({
-        'sin': sympy.sin,
-        'cos': sympy.cos,
-        'tan': sympy.tan,
-        'exp': sympy.exp,
-        'sqrt': sympy.sqrt,
-        'log': sympy.log,
-        'abs': sympy.Abs,
-        'Abs': sympy.Abs,
-        'Piecewise': sympy.Piecewise,
-    })
-
-    try:
-        pre = _convert_ternary(expr_str)
-        return parse_expr(pre, transformations=transformations, local_dict=local_dict)
-    except Exception:
-        # Fallback
-        return sympy.sympify(expr_str, locals=local_dict)
-
-
-def parse_odes_to_function(ode_json_str):
-    """
-    Parses a JSON string of ODEs and returns a dynamic update function.
-    """
-    try:
-        if isinstance(ode_json_str, str):
-            odes = json.loads(ode_json_str)
+    initial: List[str] # ["x=1", "y=0"]
+    
+    external_variables: Dict[str, float]
+    time: float
+    
+    _u_t: np.ndarray | None = None # For 2nd order or momentum
+
+    def __init__(self, equation: str = "", desc: str = "", u_shape: List[str] = None):
+        self.equation = equation
+        self.desc = desc
+        self.u = np.array([])
+        
+        if u_shape:
+            self.u_shape = u_shape
+        elif equation:
+            # Infer from equation
+            var, _, _, _ = parse_ode(equation)
+            if var and var != "u":
+                self.u_shape = [var]
+            else:
+                self.u_shape = ["u"]
         else:
-            odes = ode_json_str
-    except json.JSONDecodeError as e:
-        print(f"Failed to decode JSON from LLM: {e}")
-        return None
+             self.u_shape = ["u"]
 
-    # Define standard symbols
-    x, z, vx, vz, t = sympy.symbols('x z vx vz t')
-    
-    deriv_map = {}
-    keys = ['dx', 'dz', 'dvx', 'dvz']
+        self.initial = []
+        self.external_variables = {}
+        self.time = 0.0
+        self._u_t = None
+
+    def init_state(self, shape: tuple = (1,)):
+        """
+        Initialize the state array y (self.u).
+        For a scalar ODE ("du/dt = ..."), shape corresponds to batch size (N independent systems).
+        For a vector ODE (u_shape=["x", "v"]), self.u will be (2, N) if shape=(N,).
+        """
+        self.time = 0.0
+        
+        if not self.u_shape:
+            self.u_shape = ["u"]
+            
+        num_components = len(self.u_shape)
+        
+        # If user passed a single int, wrap it
+        if isinstance(shape, int):
+            shape = (shape,)
+            
+        if num_components > 1:
+            self.u = np.zeros((num_components, *shape), dtype=float)
+        else:
+            self.u = np.zeros(shape, dtype=float)
+            
+        self._u_t = np.zeros_like(self.u)
+
+    def set_initial_state(self):
+        """
+        Parses `self.initial` and applies to `self.u`.
+        """
+        env = self._build_eval_env()
+        env["t"] = 0.0
+        
+        if not self.initial:
+            return
+
+        for ic_eqn in self.initial:
+            if "=" not in ic_eqn: continue
+            lhs, rhs = ic_eqn.split("=", 1)
+            lhs = lhs.strip()
+            rhs_expr = rhs.strip()
+            
+            # Determine target component
+            target_idx = None
+            target_name = "u"
+            
+            m = re.match(r"^([a-zA-Z_]\w*)", lhs)
+            if m:
+                target_name = m.group(1)
+            
+            if target_name == "u" and len(self.u_shape) == 1:
+                target_idx = None
+            elif target_name in self.u_shape:
+                target_idx = self.u_shape.index(target_name)
+            
+            try:
+                rhs_expr = _preprocess_expr(rhs_expr)
+                val = eval(rhs_expr, {}, env)
+                
+                if target_idx is None:
+                    if np.shape(val) == np.shape(self.u):
+                        self.u[:] = val
+                    else:
+                        self.u[:] = val # Broadcast
+                elif len(self.u_shape) == 1 and target_idx == 0:
+                     self.u[:] = val
+                else:
+                     if target_idx < len(self.u):
+                             self.u[target_idx][:] = val
+            except Exception as e:
+                print(f"Failed to set IC '{ic_eqn}': {e}")
+
+    def _build_eval_env(self) -> Dict[str, object]:
+        
+        def pos(u): return np.maximum(u, 0.0)
+        def sign(u): return np.sign(u)
+        def step_fn(u): return np.heaviside(u, 1.0)
+
+        env = {
+            "np": np,
+            "sin": np.sin, "cos": np.cos, "tan": np.tan,
+            "sinh": np.sinh, "cosh": np.cosh, "tanh": np.tanh,
+            "arcsin": np.arcsin, "arccos": np.arccos, "arctan": np.arctan,
+            "log": np.log, "log10": np.log10, "log2": np.log2,
+            "exp": np.exp, "sqrt": np.sqrt, "abs": np.abs,
+            "pi": np.pi, "inf": np.inf, 
+            "pos": pos, "sign": sign, "step": step_fn
+        }
+        
+        env.update(self.external_variables)
+        
+        if "t" not in env:
+            env["t"] = self.time
+
+        return env
+
+    def evaluate_rhs(self, rhs_expr: str, env: Dict[str, Any]) -> np.ndarray:
+        rhs_expr = _preprocess_expr(rhs_expr)
+        return eval(rhs_expr, {}, env)
     
-    for key in keys:
-        expr_str = odes.get(key, "0")
+    def evaluate_scalar(self, expr: str, env: Dict[str, Any]) -> float:
+        if expr in ("", "1"): return 1.0
+        val = self.evaluate_rhs(expr, env)
+        if np.isscalar(val): return float(val)
+        return float(np.mean(val))
+
+    def step(self, dt: float):
+        var_name, order, coeff_expr, rhs_expr = parse_ode(self.equation)
+        
+        self.time += dt
+        env = self._build_eval_env()
+        env["t"] = self.time
+        
+        # Inject state
+        if len(self.u_shape) == 1:
+            name = self.u_shape[0]
+            env[name] = self.u
+            v_t = self._u_t if self._u_t is not None else np.zeros_like(self.u)
+            env[f"{name}_t"] = v_t
+        else:
+            for i, name in enumerate(self.u_shape):
+                env[name] = self.u[i]
+                v_t = self._u_t[i] if self._u_t is not None else np.zeros_like(self.u[i])
+                env[f"{name}_t"] = v_t
+            # Also inject 'u' as general access if needed, or maybe not to avoid confusion?
+            # PDE does env["u"] = self.u.
+            env["u"] = self.u
+            env["u_t"] = self._u_t if self._u_t is not None else np.zeros_like(self.u)
+
         try:
-            # Parse the expression safely using robust parser
-            expr = robust_parse(str(expr_str))
+            rhs = self.evaluate_rhs(rhs_expr, env)
+            if isinstance(rhs, list):
+                rhs = np.array(rhs)
+                
+            coeff = self.evaluate_scalar(coeff_expr, env)
+            forcing = rhs / (coeff if coeff else 1.0)
+            
+            # Ensure shape compatibility if forcing is lower dim (e.g. constant vector [1,2] vs shape (2, N))
+            if isinstance(forcing, np.ndarray) and forcing.ndim < self.u.ndim:
+                 diff = self.u.ndim - forcing.ndim
+                 # Assume alignment on first dimension (components)
+                 # Expand trailing dims
+                 new_shape = forcing.shape + (1,) * diff
+                 forcing = forcing.reshape(new_shape)
+
+            target_indices = []
+            if var_name == "u" and len(self.u_shape) == 1:
+                target_indices = [None]
+            elif var_name in self.u_shape:
+                idx = self.u_shape.index(var_name)
+                target_indices = [idx]
+            elif var_name == "u" and len(self.u_shape) > 1:
+                target_indices = range(len(self.u_shape))
             
-            # Create a localized function
-            # Arguments match the order we will call them
-            func = sympy.lambdify((x, z, vx, vz, t), expr, modules=['numpy', 'math'])
-            deriv_map[key] = func
+            # 2nd Order Euler / Semi-Implicit
+            if order == 2:
+                if self._u_t is None: self._u_t = np.zeros_like(self.u)
+                
+                if var_name == "u" and len(self.u_shape) > 1:
+                    self._u_t += dt * forcing
+                    self.u += dt * self._u_t
+                elif len(target_indices) == 1:
+                    idx = target_indices[0]
+                    # Handle flat u for single component
+                    is_flat = (len(self.u_shape) == 1)
+                    
+                    target_u = self.u if (idx is None or is_flat) else self.u[idx]
+                    target_ut = self._u_t if (idx is None or is_flat) else self._u_t[idx]
+                    
+                    target_ut[:] += dt * forcing
+                    target_u[:] += dt * target_ut
+                    
+            else:
+                # 1st Order Euler
+                if var_name == "u" and len(self.u_shape) > 1:
+                    self.u += dt * forcing
+                elif len(target_indices) == 1:
+                     idx = target_indices[0]
+                     if idx is None or (len(self.u_shape) == 1):
+                         self.u[:] += dt * forcing
+                     else:
+                         self.u[idx] += dt * forcing
+                         
         except Exception as e:
-            print(f"Error parsing expression for {key}: {e}")
-            return None
-
-    def dynamics(particle, dt):
-        # Current state
-        cx, cz, cvx, cvz = particle.x, particle.z, particle.vx, particle.vz
-        # We assume particle might track time, or we just pass 0 if autonomous
-        ct = getattr(particle, 'time', 0.0)
+            print(f"Error stepping ODE: {e}")
+            raise e
+
+    def get_grid(self, u_state: np.ndarray = None, dt: float = 0.0) -> np.ndarray:
+        """
+        Calculates the change (du) or rate of change (forcing/dydt) for the current state
+        without modifying the internal state.
+        
+        Args:
+            u_state: Optional state to substitute self.u
+            dt: if > 0, returns delta. If 0, returns rate.
+        """
+        original_u = self.u
+        original_u_t = self._u_t
+        
+        if u_state is not None:
+            self.u = u_state
+            if self._u_t is not None and self._u_t.shape != u_state.shape:
+                self._u_t = np.zeros_like(u_state)
+            elif self._u_t is None:
+                self._u_t = np.zeros_like(u_state)
+                
+        var_name, order, coeff_expr, rhs_expr = parse_ode(self.equation)
         
+        env = self._build_eval_env()
+        env["t"] = self.time
+
+        if len(self.u_shape) == 1:
+            name = self.u_shape[0]
+            env[name] = self.u
+            v_t = self._u_t if self._u_t is not None else np.zeros_like(self.u)
+            env[f"{name}_t"] = v_t
+        else:
+            for i, name in enumerate(self.u_shape):
+                env[name] = self.u[i]
+                v_t = self._u_t[i] if self._u_t is not None else np.zeros_like(self.u[i])
+                env[f"{name}_t"] = v_t
+            env["u"] = self.u
+            env["u_t"] = self._u_t if self._u_t is not None else np.zeros_like(self.u)
+
         try:
-            # Calculate derivatives
-            val_dx = deriv_map['dx'](cx, cz, cvx, cvz, ct)
-            val_dz = deriv_map['dz'](cx, cz, cvx, cvz, ct)
-            val_dvx = deriv_map['dvx'](cx, cz, cvx, cvz, ct)
-            val_dvz = deriv_map['dvz'](cx, cz, cvx, cvz, ct)
-            
-            # Simple Euler Integration
-            particle.x += float(val_dx) * dt
-            particle.z += float(val_dz) * dt
-            particle.vx += float(val_dvx) * dt
-            particle.vz += float(val_dvz) * dt
-            
-            # Update time if tracked
-            if hasattr(particle, 'time'):
-                particle.time += dt
+            rhs = self.evaluate_rhs(rhs_expr, env)
+            if isinstance(rhs, list):
+                rhs = np.array(rhs)
                 
+            coeff = self.evaluate_scalar(coeff_expr, env)
+            forcing = rhs / (coeff if coeff else 1.0)
+            
+            if order == 2:
+                 # 2nd Order
+                 if dt > 0:
+                     v = self._u_t if self._u_t is not None else np.zeros_like(self.u)
+                     du = dt * v + 0.5 * (dt**2) * forcing
+                 else:
+                     du = forcing
+            else:
+                 # 1st Order
+                 if dt > 0:
+                     du = forcing * dt
+                 else:
+                     du = forcing
+                     
         except Exception as e:
-            # Prevent crashing the renderer on math errors (e.g. div by zero)
-            print(f"Runtime error in dynamics: {e}")
+            self.u = original_u
+            self._u_t = original_u_t
+            raise e
+            
+        self.u = original_u
+        self._u_t = original_u_t
+        return du
+
+    def to_json(self) -> str:
+        return json.dumps({
+            "equation": self.equation,
+            "desc": self.desc,
+            "u_shape": self.u_shape,
+            "initial": self.initial,
+            "variables": self.external_variables,
+            "time": self.time
+        }, indent=2)
 
-    return dynamics
+    def __str__(self):
+        return self.to_json()