stepping.py

import numpy as np
import scipy.integrate as itg


def solve_discrete_difference_eqns(discrete_state_model, x0, u_fcn, DT, T):
    """

    :param discrete_state_model:
    :param x0: m, |
    :param u_fcn: (x_t, t) -> u_t \in n, |
    :param DT: sampling time
    :param T: final time, steps made for t < T
    :return:
        x: m, k | ith column is state at t = i*DT
        u: n, k | ith column is control at t
    """
    m, n = discrete_state_model.get_mn()
    t = np.arange(0, T, DT)
    x = np.zeros((m, t.size))
    u = np.zeros((n, t.size))
    x[:, 0] = x0.copy()
    for i in range(0, t.size-1):
        discrete_state_model.update_AB(x[:, i])
        u[:, i] = u_fcn(x[:, i], i*DT)
        x[:, i+1] = discrete_state_model.step(x[:, i], u[:, i])
    return x, u


def solve_cts_eqns(cts_state_model, x0, u_fcn, DT, T):
    """
    :param cts_state_model:
    :param x0: m, |
    :param u_fcn: (x_t, t) -> u_t \in n, |
    :param DT: sampling time
    :param T: final time, steps made for t < T
    :return:
        x: m, k | ith column is state at t = i*DT
        u: n, k | ith column is control at t
    """
    m, n = cts_state_model.get_mn()
    t = np.arange(0, T, DT)
    x = np.zeros((m, t.size))
    u = np.zeros((n, t.size))
    x[:, 0] = x0.copy()
    for i in range(0, t.size - 1):
        cts_state_model.update_AB(x[:, i])
        A, B = cts_state_model.get_AB()
        u[:, i] = u_fcn(x[:, i], i * DT, A=A, B=B)
        t_span = (i*DT, (i+1 + 0.5)*DT)
        t_eval = (i+1)*DT,
        sol = itg.solve_ivp(
            lambda t, x: cts_state_model.step(x, u[:, i]),
            t_span, x[:, i], t_eval=t_eval)
        x[:, [i + 1]] = sol.y
    return x, u