from sympy import ( Abs, Add, Basic, Dummy, Eq, Function, Matrix, re, solve, Subs, Symbol, symbols, Tuple, ) from typing import Any from sympy.core.function import AppliedUndef from sympy.logic import And from sympy.integrals import laplace_transform, LaplaceTransform def symbol_matrix(name, length): return Matrix(symbols(f"{name}_:{length}", real=True)) def function_matrix(name, length, arg): return Matrix( [func(arg) for func in symbols(f"{name}_:{length}", cls=Function, real=True)] ) def _everything(ode, discretize, time, sample_freq, n, point): fde = FDE.from_ode(ode, discretize, time, sample_freq) solved_fde = SolvedFDE.from_fde(fde, n) return ( ode.stable(point), fde.stable(point), solved_fde.stable(point), solved_fde, ) class ODE(Basic): # Assumes state-space form: # # d # ----- state = expr(state, time) # dtime # # See https://ccrma.stanford.edu/~jos/pasp/General_Nonlinear_ODE.html # def __init__(self, time: Symbol, state: Matrix, expr: Matrix): self.time = time self.state = state self.expr = expr def stable(self, operating_point: Matrix): return _ode_stability_condition( self.expr, self.state, self.time, operating_point ) class FDE(Basic): # Assumes the form # # 0 = expr(state, time) # def __init__(self, time: Symbol, state: Matrix, expr: Matrix): self.time = time self.state = state self.expr = expr @classmethod def from_ode(cls, ode: ODE, discretize, time: Symbol, sample_freq: Symbol): return FDE( time, ode.state.subs(ode.time, time), discretize(ode.state, ode.expr, ode.time, time, sample_freq), ) def stable(self, operating_point: Matrix): return _fde_stability_condition( self.expr, self.state, self.time, operating_point ) class SolvedFDE(Basic): # Assumes the form # # state = expr(state, time) # # which must be an explicit equation. # def __init__(self, time: Symbol, state: Matrix, expr: Matrix): self.time = time self.state = state self.expr = expr @classmethod def from_fde(cls, fde: FDE, n): return SolvedFDE( fde.time, fde.state, _solve_fde(fde.state, fde.expr, fde.time, n), ) def stable(self, operating_point: Matrix): return _fde_stability_condition( self.expr, self.state, self.time, operating_point ) def backward_euler(y, f, t, k, f_s): return Matrix( [ -f_s * _backward_difference(y[i], t, k) + f[i].subs(t, k) for i in range(f.shape[0]) ] ) def trapezoidal(y, f, t, k, f_s): return Matrix( [ -2 * f_s * _backward_difference(y[i], t, k) + f[i].subs(t, k) + f[i].subs(t, k - 1) for i in range(f.shape[0]) ] ) def _backward_difference(y, t, k): return y.subs(t, k) - y.subs(t, k - 1) def _newton(n, y, func, k): guess = y.subs(k, k - 1) return Newton(n, func, y, guess) class Newton(Function): nargs = 4 @classmethod def eval(cls, n, f, var, guess): if not n.is_Number: return None if n == 0: return guess if n == 1: return guess - (f / f.diff(var)).subs(var, guess) if n > 1: expr = guess for _ in range(n): expr = Newton(1, f, var, expr) return expr def _solve_fde(y, f, k, n): def solution(i): try: return solve(f[i], y[i])[0] except (IndexError, NotImplementedError): return _newton(n[i], y[i], f[i], k) return Matrix([solution(i) for i in range(f.shape[0])]) # stability checking # ↓ def _linearize(func, state, point): """ Scalar multi-variate function #x ∂ f(x) ≅ Σ f(p[n]) + ----- f | (x[n] - p[n]) n=0 ∂x[n] |x=p Vector multi-variate function #x ∂ f[m](x) ≅ Σ f[m](p[n]) + ----- f[m] | (x[n] - p[n]) n=0 ∂x[n] |x=p """ substs = [(state[n], point[n]) for n in range(func.shape[0])] return Matrix( [ Add( func[m].subs(substs), *[ func[m].diff(state[n]).subs(substs) * (state[n] - point[n]) for n in range(state.shape[0]) ], ) for m in range(func.shape[0]) ] ) def _block_dc(func, time): return Matrix( [ eqn.func(*list(filter(lambda term: term.has(time), eqn.args))) for eqn in func.expand() ] ) def _reduce(state, func, time): func = list(func) targets = state[1:] while len(targets) > 0 and len(func) > 1: target = targets.pop() index, solution = _index_solution(func, target, time) func = _remove(func, index).subs(target, solution) assert len(targets) == 0 assert len(func) == 1 return func[0].simplify() def _index_solution(func, target, time): for index, eqn in enumerate(func): if eqn.has(target.diff(time)): continue solutions = solve(eqn, target) if len(solutions) == 1: return index, solutions[0] assert len(solutions) == 0 return None, None def _remove(func, index): return Matrix(func[:index] + func[index + 1 :]) def _upper_symbol(expr): return Symbol(str(expr.func).upper()) def _laplace_transform(char_eqn, t, s): def _symbolized_laplace(lt): expr = lt.args[0] if expr.is_Derivative: return s ** expr.args[1][1] * _upper_symbol(expr.args[0]) elif isinstance(expr, AppliedUndef): return _upper_symbol(expr) raise BaseException(f"unexpected laplace transform of {expr}") lt_result = laplace_transform(char_eqn, t, s) if hasattr(lt_result, "__getitem__"): raise BaseException(char_eqn) return lt_result.replace( lambda e: e.func == LaplaceTransform, _symbolized_laplace, ) def _ode_stable(char_eqn, t): s = Dummy("s") return And(*[re(root) < 0 for root in solve(_laplace_transform(char_eqn, t, s), s)]) def _ode_stability_condition(func, state, time, point): return _ode_stable( _reduce( state, _block_dc(_linearize(func, state, point), time) - state.diff(time), time, ), time, ) # Try this version when trying to eliminate .has(state.diff(time)) check in _reduce. # return _ode_stable(_reduce(state, _block_dc(_linearize(func, state, point), time), time) - state.diff(time), time) def _get_shift(signal, k): time = signal.args[0] d = Dummy("d") return solve(Eq(time, k + d), d)[0] def _z_transform(expr, k, z): return expr.replace( lambda e: isinstance(e, AppliedUndef) and len(e.args) == 1 and e.has(k), lambda e: z ** _get_shift(e, k) * _upper_symbol(e), ) def _fde_stable(y_k, f_k_, k): z = Dummy("z") Y, F = _z_transform(Tuple(y_k, f_k_), k, z) return And(*[Abs(root) < 1 for root in solve(_reduce(Y, F, Dummy()), z)]) def _fde_stability_condition(f_k, y_k, k, point): return _fde_stable(y_k, _block_dc(_linearize(f_k, y_k, point) - y_k, k), k)