implement visualize_trajectory_3d

jobs/src/visualization.py

+import matplotlib.pyplot as plt
+import inspect
+
+def visualize_trajectory_3d(trajectory):
+    '''
+    input:
+    trajectory should be 3D nparray:
+    first: the number of particles
+    second: 3, the coordinates of x,y,z
+    third: the nparray with the capacity of the number of steps
+    '''
+
+    print(f"visualize the trajectory_3d:--------------Running---------------------")
+    # figure
+    plt.style.use('dark_background')
+    fig = plt.figure()
+    ax = fig.add_subplot(projection='3d')
+
+    # # choose a different color for each trajectory
+    # colors = plt.cm.jet(np.linspace(0, 1, N))
+    #
+    # # set up lines and points
+    # lines = sum([ax.plot([], [], [], '-', c=c)
+    #              for c in colors], [])
+    # pts = sum([ax.plot([], [], [], 'o', c=c)
+    #            for c in colors], [])
+    steps = len(trajectory[0][0])
+    for i in range(steps):
+        # print (i)
+        # clear the current figure
+        plt.cla()
+        xp = trajectory[:, 0, max(i - 50, 0):i + 1]
+        yp = trajectory[:, 1, max(i - 50, 0):i + 1]
+        zp = trajectory[:, 2, max(i - 50, 0):i + 1]
+
+
+        # plotting the trails
+        ax.scatter(xp, yp, zp, s=1, color='red')
+        ax.scatter(trajectory[:, 0, i], trajectory[:, 1, i], trajectory[:, 2, i], s=10, color='white')
+        ax.set_xlabel('X Label')
+        ax.set_ylabel('Y Label')
+        ax.set_zlabel('Z Label')
+        ax.set_title(f'step:{i}, Total Step:{steps}')
+        plt.pause(0.001)
+    plt.show()