File: ode2.py

package info (click to toggle)
python-bumps 1.0.1-1
  • links: PTS, VCS
  • area: main
  • in suites: experimental
  • size: 6,180 kB
  • sloc: python: 24,284; xml: 493; ansic: 373; makefile: 209; sh: 94; javascript: 88
file content (163 lines) | stat: -rw-r--r-- 4,116 bytes parent folder | download | duplicates (2)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
# 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)