Make the animation script accept any (.mat) file

diff --git a/animate_result.py b/animate_result.py
index 2052ef4f585bce21d0b72e9d5cac5ba65b403369..1e8fbc8d5670de6a9b9147cd451ff9e9bf5c0ec6 100755 (executable)
--- a/animate_result.py
+++ b/animate_result.py
@@ -1,5 +1,6 @@
 #!/usr/bin/env python3
 
+import sys
 import matplotlib.pyplot as plt
 import numpy as np
 import matplotlib.animation as animation
@@ -57,15 +58,19 @@ def update(frame):
     line.set_xdata(datadict[name_x][frame]);
     line.set_ydata(datadict[name_y][frame]);
 
+def matAniFromFile(filename):
+    data = scp.io.loadmat("nbody.System_res.mat");
+    names = transpose_strings(data["name"]);
+    descriptions = transpose_strings(data["description"]);
+    pure_data = data["data_2"];
+    inds = data["dataInfo"][1];
+    ani = MatAnimation();
+    for i in inds:
+        ani.data[names[i - 1]] = pure_data[i - 1, :-1];
+    return (ani);
 
-data = scp.io.loadmat("nbody.System_res.mat");
-names = transpose_strings(data["name"]);
-descriptions = transpose_strings(data["description"]);
-pure_data = data["data_2"];
-inds = data["dataInfo"][1];
-ani = MatAnimation();
-for i in inds:
-    ani.data[names[i - 1]] = pure_data[i - 1, :-1];
+
+ani = matAniFromFile(sys.argv[1]);
 ani.plot("p[1].pos[1]", "p[1].pos[2]");
 ani.plot("p[2].pos[1]", "p[2].pos[2]");
 ani.plot("p[3].pos[1]", "p[3].pos[2]");