In [ ]:
import scipy.integrate
import scipy.sparse
import numpy as np
import collections
import matplotlib.pyplot as plt
%matplotlib inline

class IVPResult:
    pass

def solve_ivp(f, ts, x0, p=None, integrator='dopri5', store_trajectory=False,
              sens=False, dfx=None, dfp=None):
    """
    Solve initial value problem
        d/dt x = f(t, x, p);  x(t0) = x0.

    Evaluate Solution at time points specified in ts, with ts[0] = t0.

    Compute sensitivity matrices evaluated at time points in ts if sens=True.
    """

    if sens:
        def f_vode(t, xGxGp, p):
            x     = xGxGp[:x0.shape[0]]
            Gx    = xGxGp[x0.shape[0]:x0.shape[0]+x0.shape[0]*x0.shape[0]]
            Gx    = Gx.reshape([x0.shape[0], x0.shape[0]])
            Gp    = xGxGp[x0.shape[0]+x0.shape[0]*x0.shape[0]:]
            Gp    = Gp.reshape([x0.shape[0], p.shape[0]])
            dx    = f(t, x, p)
            dfxev = dfx(t, x, p)
            dfpev = dfp(t, x, p)
            dGx   = dfxev.dot(Gx)
            dGp   = dfxev.dot(Gp) + dfpev
            return np.concatenate(
                [dx,
                 dGx.reshape(-1),
                 dGp.reshape(-1)])

        ivp = scipy.integrate.ode(f_vode)
    else:
        ivp = scipy.integrate.ode(f)

    ivp.set_integrator(integrator)

    if store_trajectory:
        times = []
        points = []
        def solout(t, x):
            if len(times) == 0 or t != times[-1]:
                times.append(t)
                points.append(np.copy(x[:x0.shape[0]]))
        ivp.set_solout(solout)

    if sens:
        ivp.set_initial_value(np.concatenate(
            [x0,
             np.eye(x0.shape[0]).reshape(-1),
             np.zeros([x0.shape[0], p.shape[0]]).reshape(-1)
            ]), ts[0])
    else:
        ivp.set_initial_value(x0, ts[0])
    ivp.set_f_params(p)

    result = IVPResult()
    result.ts = ts
    result.xs  = np.zeros([ts.shape[0], x0.shape[0]])
    if sens:
        result.Gxs = np.zeros([ts.shape[0], x0.shape[0], x0.shape[0]])
        result.Gps = np.zeros([ts.shape[0], x0.shape[0], p.shape[0]])
    result.success = True

    result.xs[0,:] = x0
    for ii in range(1,ts.shape[0]):
        ivp.integrate(ts[ii])
        result.xs[ii,:]    = ivp.y[:x0.shape[0]]
        if sens:
            result.Gxs[ii,:,:] = ivp.y[x0.shape[0]:x0.shape[0]+x0.shape[0]*x0.shape[0]].reshape([x0.shape[0], x0.shape[0]])
            result.Gps[ii,:,:] = ivp.y[x0.shape[0]+x0.shape[0]*x0.shape[0]:].reshape([x0.shape[0], p.shape[0]])
        if not ivp.successful():
            result.success = False
            break

    if store_trajectory:
        result.trajectory_t = np.array(times)
        result.trajectory_x = np.array(points)

    return result


# delta x should be a vector of dimension equal to the number
# of variables, which is true only if we do not change the parameters.
def evaluate_parest_single_shooting(f, dfx, dfp, meas_times, meas_values, s0, p):
    result  = solve_ivp(f, meas_times, s0, p=p, sens=True, dfx=dfx, dfp=dfp)
    m = len(meas_times)
    nx = len(s0)
    nv = nx + len(p)
    # make F a column vector with [y11, y12, y21, y22, y31, y32, ... ym1, ym2]
    F = (meas_values - result.xs).reshape((m * nx, 1))
    # each row of J will be [dFij / dx01, dFij / dx02, dFij / dp1, ..., dFij / dp6 ]
    J = np.zeros((m * nx, nv))
    dy = np.eye(nx)
    for i in range(m):
        for j in range(nx):
            # dhi / dy = dyi / dy = all zeros except for the ith element
            # that will be 1
            # For dFij / dp I  dhi / dp = dyi(tj, yi, p) / dp is equal to 0
            J[2 * i + j, :] = np.hstack((dy[[j], :].dot(result.Gxs[i]), dy[[j], :].dot(result.Gps[i])))
    return F, J


def gauss_newton(F_J, x0, meas_times, meas_values, itmax=100, tol=1e-7, verbose=1):
    
    s0 = x0[:2]
    param = x0[2:]
    tau = .5
    if verbose == 1:
        print("Start gauss_newton with itmax=",itmax,"and tol=",tol,"and tau=",tau)
        Fs = np.zeros(itmax+1)
    i = 0
    norm = tol+1.0
    while i<itmax and norm > tol:
        F, J= F_J(dX_dt, dfx, dfp, meas_times, meas_values, s0, param)
        Dx = np.linalg.lstsq(J, F)[0]
        s0 += tau * Dx.flatten()[:2]
        param += tau * Dx.flatten()[2:]
        norm = np.linalg.norm(Dx.flatten())
        if verbose == 1:
            Fs[i] = (np.linalg.norm(F))
            print("It.step: ",i+1, '\t', u'\u2016Dx\u2016\u2082 =', norm)
        i += 1
        
            
    if verbose == 1:
        print("\n The shooting residual F throughout the iterations:")
        plt.plot(Fs[:i])
        Fs
        
    return s0, param