stepping.py

import numpy as np


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