python-lib/animate_xz.py at master · AquaDragon/python-lib · GitHub

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
'''
NAME:           animate_xz.py
AUTHOR:         swjtang
DATE:           15 Jan 2021
DESCRIPTION:    Plots an animated contour plot on the XZ plane
SYNTAX:
'''
import numpy as np
from matplotlib import animation, cm, pyplot as plt

from lib.read_hdf5_v2 import readhdf5


def animate(fname, dirw='', cmap=cm.GnBu):
    dataset = readhdf5(dirw+fname)

    pos = dataset['pos']
    time = dataset['time']
    bx = dataset['ch1']
    by = dataset['ch2']
    bz = dataset['ch3']

    nx = dataset['nx']
    ny = dataset['ny']
    nt = len(time)

    xx = dataset['xx']
    yy = dataset['yy']
    yy = 20*np.sin(np.deg2rad(yy))  # list of y positions

    btot = [np.sqrt(bx[i]**2 + by[i]**2 + bz[i]**2) for i in range(len(bx))]
    # reshape is equivalent to reform
    data = np.reshape(btot, (ny, nx, nt), order='C')

    fig = plt.figure(figsize=(15, 10))  # , constrained_layout=True)
    fig.suptitle(str(fname), fontsize=20)
    ax = fig.add_subplot(331)
    ay = fig.add_subplot(332)
    az = fig.add_subplot(333)
    atot = fig.add_subplot(312, rowspan=2)

    vmin = 0.6*np.amin(data)
    vmax = 0.6*np.amax(data)
    levels = np.linspace(vmin, vmax, 30, endpoint=True)

    frame = atot.contourf(xx, yy, data[:, :, 0], vmax=vmax, vmin=vmin,
                          levels=levels, cmap=cmap)
    cbar = fig.colorbar(frame, ax=atot)
    cbar.ax.tick_params(labelsize=16)
    cbar.set_label('Magnitude', fontsize=16)

    # animation parameters
    nstep = 100                      # frames are taken at every dt step size
    nframes = len(time)//nstep

    # animation function
    def generate_frame(i):
        atot.clear()
        atot.contourf(xx, yy, data[:, :, nstep*i], vmax=vmax, vmin=vmin,
                      levels=levels, cmap=cmap)
        atot.set_title('t = {0:.4f}  us [{1}]'.format(time[nstep*i]*1e6, i),
                       fontsize=18)
        atot.set_xlabel('X [cm]', fontsize=18)
        atot.set_ylabel('Y [cm]', fontsize=18)
        atot.tick_params(labelsize=16)
        atot.set_aspect('equal')

        ax.contourf(xx, yy, data[:, :, nstep*i], vmax=vmax, vmin=vmin,
                    levels=levels, cmap=cmap)
        ax.set_aspect('equal')

        ay.contourf(xx, yy, data[:, :, nstep*i], vmax=vmax, vmin=vmin,
                    levels=levels, cmap=cmap)
        ay.set_aspect('equal')

        az.contourf(xx, yy, data[:, :, nstep*i], vmax=vmax, vmin=vmin,
                    levels=levels, cmap=cmap)
        az.set_aspect('equal')

        print('\r', 'Generating frame {0}/{1} ({2:.2f}%)...'.
              format(i+1, nframes, (i+1)/nframes*100), end='')

    anim = animation.FuncAnimation(fig, generate_frame, interval=50,
                                   frames=nframes)

    temp = fname.find('.')
    anim.save(fname[0:temp]+'.mp4')
    plt.show()