Secrets of Spring Motion - Figure 5

fig5.py

# Figure 5 from http://fouryears.eu/2016/11/17/the-secrets-of-spring-motion/
# Intended to be run from a Jupyter notebook cell

%pylab inline

g = 9.18
k = 2.0 * 8
L = 1.0
m = 1.0/9

class Model(object):
    def __init__(self, n):
        self.ms = ones(n)*m
        self.vs = zeros(n)
        # Initially balanced xs
        self.xs = concatenate([[0],-cumsum(cumsum(self.ms[::-1])[::-1][1:]*g/k + L)])
        assert all(abs(self.fs[1:]) < 1e-10)

    @property
    def fs(self):
        ds = abs(self.xs[1:] - self.xs[:-1])
        f_stretch = (ds - L)*k
        f_stretch_up = concatenate([[0], f_stretch])
        f_stretch_down = concatenate([-f_stretch, [0]])
        f_g = -self.ms*g
        fs = f_stretch_up + f_stretch_down + f_g
        return fs

    def step(self, dt):
        self.vs += self.fs/self.ms*dt
        self.xs += self.vs*dt
    
    def simulate(self, sim_time=1.0, n_steps=10000):
        x_trace = []
        f_trace = []
        dt = sim_time / n_steps
        for i in range(n_steps):
            x_trace.append(self.xs.copy())
            f_trace.append(self.fs.copy())
            self.step(dt)
        return arange(n_steps)*dt, np.asarray(x_trace), np.asarray(f_trace)


model = Model(10)
ts, xs, fs = model.simulate()
figure(figsize=(6,4))
plot(ts, xs)
grid()