How to make an animation of rocket motion in matpl

Question

I have this code to calculate position, velocity and acceleration of a certain rocket according to given data:

import math
import matplotlib.pyplot as plt
import pylab as pl
import numpy as np

M = 30*10**3
m = 10*10**3
fi = 100
vr = 3*10**3
x = np.linspace (0, 324, 1296)

def v(i):
    if i <= 225:
        y = vr*math.log((M+m)/(m+M - fi*i))
    elif i > 225:
        y = vr*math.log((M+m)/(m+M/4)) + vr*math.log((m)/(m - fi*(i - 225)))
    return y

def a(i):
    if i <= 225:
        y = fi*vr/(m+M - fi*i)
    elif i > 225:
        y = fi*vr/(m - fi*(i - 225))
    return y

def z(i):
    if i <= 225:
        y = vr/fi*(m+M-fi*i)*(math.log((1-fi*i/(m+M)))-1) + vr*(m+M)/fi
    elif i > 225:
        y = vr*(i-225)*math.log((m+M)/(m+M/4)) + z(225) + vr/fi*(m-fi*(i-225))*(math.log((1-fi*(i-225)/(m)))-1) + vr*m/fi
    return y


plt.figure(0)
ax1 = plt.subplot2grid((4,3), (0,1), colspan=2, rowspan=2)
plt.grid()
plt.plot (x, [v(i) for i in x], 'b--') 
plt.ylabel (r'Velocity [$\frac{m}{s}$]', fontsize = 16)
plt.xlabel ('Time [s]', fontsize = 16)
plt.title ('Velocity with respect to time', fontsize = 18)

ax2 = plt.subplot2grid((4,3), (0,0), colspan=1, rowspan=4)
plt.grid()
plt.plot (x, [z(i) for i in x], 'b--')
plt.ylabel (r'Position [m]', fontsize = 16)
plt.xlabel ('Time [s]', fontsize = 16)
plt.title ('Position with respect to time', fontsize = 18)

ax3 = plt.subplot2grid((4,3), (2,1), rowspan=2, colspan=2)
plt.grid()
plt.plot (x, [a(i) for i in x], 'b--') 
plt.xlabel ('Time [s]', fontsize = 16)
plt.ylabel (r'Acceleration [$\frac{m}{s^2}$]', fontsize = 16)
plt.title ('Acceleration with respect to time', fontsize = 18)

plt.suptitle ('Rocket behaviour', fontweight = 'bold', fontsize = 20)
plt.show()

But, instead of the 'Position with respect to time' graph, I want to make an animation of the rocket motion in the z-direction (vertically upwards). How could I do that? All my attempts have failed, so I kindly ask for some assistance. The animation should be in animated GIF format.

Answer 1

This shows an animation:

import math

import matplotlib.pyplot as plt
import numpy as np
import matplotlib.animation as animation


def z(i):
    if i <= 225:
        y = (vr / fi * (m + M - fi * i) *
            (math.log((1 - fi * i / (m + M))) - 1)
            + vr * (m + M) / fi)
    elif i > 225:
        y = (vr * (i - 225) * math.log((m + M) / (m + M / 4)) + z(225) +
             vr / fi * (m - fi * (i - 225)) *
             (math.log((1 - fi * (i - 225) / (m))) - 1) + vr * m / fi)
    return y


class GrowingGraph(object):
    """The z and time axis are growing over time.
    """

    def __init__(self):
        self.x_data = []
        self.y_data = []
        self.fig, self.ax = plt.subplots()
        plt.ylabel(r'Position [m]', fontsize=16)
        plt.xlabel('Time [s]', fontsize=16)
        plt.title('Position with respect to time', fontsize=18)

    def run(self, data):
        # update the data
        x, y = data
        self.x_data.append(x)
        self.y_data.append(y)
        line, = self.ax.plot(self.x_data, self.y_data, color='blue')
        return line,

class FixedGraph(object):
    """z and time axis are fixed at maximum values.
    """

    def __init__(self, xmax, ymax):
        self.x_data = []
        self.y_data = []
        self.fig = plt.figure()
        self.ax = plt.axes(xlim=(0, xmax), ylim=(0, ymax))
        self.line, = self.ax.plot(self.x_data, self.y_data)
        self.line.set_data([], [])
        plt.ylabel(r'Position [m]', fontsize=16)
        plt.xlabel('Time [s]', fontsize=16)
        plt.title('Position with respect to time', fontsize=18)


    def run(self, data):
        # update the data
        x, y = data
        self.x_data.append(x)
        self.y_data.append(y)
        self.line.set_data(self.x_data, self.y_data)
        return self.line,

if __name__ == '__main__':

    M = 30 * 10 ** 3
    m = 10 * 10 ** 3
    fi = 100
    vr = 3 * 10 ** 3
    x = np.linspace(0, 324, 1296)

    data = zip(x, (z(i) for i in x))

    graph = FixedGraph(x[-1], z(x[-1]))
    ani = animation.FuncAnimation(graph.fig, graph.run, frames=data,
                                  interval=200, repeat=False)

    # Comment out this line if you don't want see the animation.
    plt.show()

    # save as gif
    ani.save('rocket.gif', writer='imagemagick', fps=4)

    # save as mp4
    ani.save('rocket.mp4', writer='ffmpeg', fps=15)

The class FuncAnimation() needs to be called with a figure instance, a callable that updates the plot and the data. The docs have the details.