update some billiards features

Author: duguyue100

src/environments/billiards.py

@@ -40,7 +40,8 @@ def __init__(self, T=128, n=2, r=None, m=None, eps=0.5, SIZE=10):
 
         self.X = np.zeros((self.T, self.n, 2), dtype=np.float)
         self.v = np.random.randn(self.n, 2)
-        self.v /= self.norm(self.v)*0.5
+        self.vnorm = self.norm(self.v)
+        self.v /= self.vnorm*0.5
 
         self.x = self.config()
 

src/run_billiards_example.py

@@ -1,6 +1,8 @@
 """Billiards Example."""
 from __future__ import print_function
 
+import numpy as np
+
 import matplotlib.pyplot as plt
 from matplotlib.animation import FuncAnimation
 
@@ -9,18 +11,20 @@
 resolution = 128
 time_steps = 25
 
-env = Billiards(T=time_steps)
+env = Billiards(T=time_steps, SIZE=10, r=np.array([0.6]*2))
 env.make_frames(resolution)
 
 frames = env.frames
 num_frames = frames.shape[0]
 
 fig = plt.figure()
 
+
 def update(t):
     plt.imshow(frames[t], cmap="gray")
 
-if __name__=="__main__":
+
+if __name__ == "__main__":
     # FuncAnimation will call the 'update' function for each frame
     anim = FuncAnimation(fig, update, frames=time_steps, interval=200)
-    anim.save('billiards.gif',writer='imagemagick');
+    anim.save('billiards-test.gif', writer='imagemagick')