From 10299ae9a188bb9e2c11d823460478f2ed67505c Mon Sep 17 00:00:00 2001
From: JayM0826 <maji91fan@gmail.com>
Date: Tue, 13 Feb 2024 11:06:08 +0100
Subject: [PATCH] implement visualize_trajectory_3d
---
jobs/src/visualization.py | 45 +++++++++++++++++++++++++++++++++++++++
1 file changed, 45 insertions(+)
diff --git a/jobs/src/visualization.py b/jobs/src/visualization.py
index e69de29..e994755 100644
--- a/jobs/src/visualization.py
+++ b/jobs/src/visualization.py
@@ -0,0 +1,45 @@
+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()
--
GitLab