Fitting a multi-valued function
===============================

Like the ODE fit function, but this example fits a set of coupled ODEs.
In this case, there are multiple values reported at each time step, two
of which are measured and fitted.

From SciPy cookbook coupled spring mass example:

   https://scipy-cookbook.readthedocs.io/items/CoupledSpringMassSystem.html


::

    from bumps.names import *
    from scipy.integrate import odeint
    
Use ODEINT to solve the differential equations defined by the vector field

::

    def vectorfield(w, t, p):
        """
        Defines the differential equations for the coupled spring-mass system.
    
        Arguments:
            w :  vector of the state variables:
                      w = [x1,y1,x2,y2]
            t :  time
            p :  vector of the parameters:
                      p = [m1,m2,k1,k2,L1,L2,b1,b2]
        """
        x1, y1, x2, y2 = w
        m1, m2, k1, k2, L1, L2, b1, b2 = p
    
        # Create f = (x1',y1',x2',y2'):
        f = [y1,
             (-b1 * y1 - k1 * (x1 - L1) + k2 * (x2 - x1 - L2)) / m1,
             y2,
             (-b2 * y2 - k2 * (x2 - x1 - L2)) / m2]
        return f
    
ODE solver parameters

::

    abserr = 1.0e-8
    relerr = 1.0e-6
    
Curve function with all parameters exposed so that bumps knows their names.
Only tracking x1, x2 with our measurements and not y1, y2, so returning
components 0 and 2 of the *vectorfield* result.  The multi-valued y values
are stacked into an array whose first axis matches t.  This is needed so that
the plotter can sort out the different lines.

::

    def f(t, x1, y1, x2, y2, m1, m2, k1, k2, L1, L2, b1, b2):
        # Pack up the parameters and initial conditions:
        p = [m1, m2, k1, k2, L1, L2, b1, b2]
        w0 = [x1, y1, x2, y2]
    
        # Call the ODE solver.
        wsol = odeint(vectorfield, w0, t, args=(p,),
                      atol=abserr, rtol=relerr)
        return np.vstack((wsol[:, 0], wsol[:, 2]))
    
Simulation parameter values

::

    # Masses
    m1 = 1.0
    m2 = 1.5
    
    # Spring constants
    k1 = 8.0
    k2 = 40.0
    
    # Natural lengths
    L1 = 0.5
    L2 = 1.0
    
    # Friction coefficients
    b1 = 0.8
    b2 = 0.5
    
Initial conditions

::

    # x1 and x2 are the initial displacements; y1 and y2 are the initial velocities
    x1 = 0.5
    y1 = 0.0
    x2 = 2.25
    y2 = 0.0
    
Simulate data

::

    def simulate():
        from bumps.util import push_seed
    
        # Create the time samples for the output of the ODE solver.
        # These are the times that the data is sampled, not the times at
        # which to evaluate the ode solver.
        t = np.linspace(0, 10, 100)
    
        # Pack up the parameters and initial conditions:
        p = [m1, m2, k1, k2, L1, L2, b1, b2]
        w0 = [x1, y1, x2, y2]
        ft = f(t, *(w0 + p))
    
        noise = 0.1*np.ones_like(ft)
        with push_seed(1):  # Make sure that the simulated data is the same each run
            data = ft + noise*np.random.randn(*ft.shape)
        return t, data, noise
    
    t, y, dy = simulate()
    
Initial values for most parameters are known from system configuration.
We are not including the spring constants or the friction coefficients
since these will be fitted to the measured position over time.  *labels*
allow you to set the labels for the x-axis and y-axis and the legend for
the two data lines on the plot.

::

    M = Curve(f, t, y, dy, m1=m1, m2=m2, L1=L1, L2=L2, x1=x1, y1=y1, x2=x2, y2=y2,
              labels=['time', 'value', 'x1', 'x2'], plot_x=np.linspace(0, 10, 1000))
    
Fitted parameters

Only fitting spring constants and friction coefficients since these are
not immediately measurable.  If we wanted to be fancy, we could set the
prior on position and mass according to the uncertainty in our
initial configuration and allow them to vary slightly.

::

    # Masses: Allow mass estimate to be off by +/- 2% (1-sigma)  *untested*
    #M.m1.dev(0.02*m1)
    #M.m2.dev(0.02*m2)
    
    # Spring constants
    M.k1.range(0, 100)
    M.k2.range(0, 100)
    
    # Natural lengths
    #M.L1.range(0, 10)
    #M.L2.range(0, 10)
    
    # Friction coefficients
    M.b1.range(0, 1)
    M.b2.range(0, 1)
    
    # Initial conditions
    # x1 and x2 are the initial displacements; y1 and y2 are the initial velocities
    #M.x1.range(0, 10)
    #M.x2.range(0, 10)
    #M.y1.range(0, 10)
    #M.y2.range(0, 10)
    
    problem = FitProblem(M)


.. only:: html

   Download: :download:`ode2.py <ode2.py>`.
