imlement saving to files

torus.py

@@ -4,6 +4,9 @@ import matplotlib.pyplot as plt
 import matplotlib.animation as animation
 from matplotlib.widgets import Slider, Button
 
+RAYTRACING = False
+DISTANCE = False
+
 
 def uproot(arr):
     mask = np.logical_and(arr > 0, np.isreal(arr))
@@ -67,7 +70,9 @@ class TorusWorld:
         rx, ry, rz = self.surface_point(phi, theta)
         sx, sy, sz = self.sun
         retx, rety, retz = rx-sx, ry-sy, rz-sz
-        return np.array([retx, rety, retz])/np.linalg.norm([retx, rety, retz])
+        length = np.linalg.norm([retx, rety, retz])
+        norm = length*length if DISTANCE else length
+        return np.array([retx, rety, retz])/norm
 
     def is_illuminated(self, phi, theta):
         ray = self.ray_vector(phi, theta)
@@ -78,13 +83,14 @@ class TorusWorld:
         return np.allclose(self.surface_point(phi, theta), self.sun+roots[0]*ray)
 
     def illumination(self, phi, theta):
-        return np.dot(self.ray_vector(phi, theta), -self.normal_vector(phi, theta))  # * self.is_illuminated(phi, theta)
+        rt = self.is_illuminated(phi, theta) if RAYTRACING else 1
+        return np.dot(self.ray_vector(phi, theta), -self.normal_vector(phi, theta)) * rt
 
 
-class NoInamge():
+class Image():
 
     def __init__(self, rfrac_init=0.5, sun_init=np.pi/2):
-        self.fig, self.ax = plt.subplots()
+        self.fig, self.ax = plt.subplots(figsize=(10, 16))
 
         self.lines = dict()
 
@@ -102,10 +108,16 @@ class NoInamge():
         self.ax.set_aspect('equal')
         self.ax.axis('off')
 
+        self.ax.set_xlim(-2.05, 2.05)
         self.init_top_view()
         self.init_side_view()
         self.init_map_view()
 
+    def save(self, path="./torus.png", dpi=72):
+        self.fig.tight_layout()
+
+        plt.savefig(path, dpi=dpi)
+
     def update_illumination(self):
         # TODO: refactor TorusWorld.illumination() do use vectoization!!
         phi = np.linspace(-np.pi, np.pi, self.res)
@@ -136,10 +148,10 @@ class NoInamge():
             ]
 
     def _offset_map(self):
-        return -self.torus.r_maj*1.1 - (1+np.pi*self._scale_map())*self.torus.r_min
+        return -self.torus.r_maj*1.3 - (1+np.pi*self._scale_map())*self.torus.r_min
 
     def _offset_side(self):
-        return self.torus.r_maj*1.1
+        return self.torus.r_maj*1.2
 
     def _scale_map(self):
         if self.rectangular_map:
@@ -301,19 +313,11 @@ class NoInamge():
         self.fig.canvas.draw_idle()
 
 
-class ImageStatic(NoInamge):
+class InteractiveImage(Image):
-
-    def __init__(self, rfrac_init=0.5, sun_init=np.pi/2):
-        super().__init__(rfrac_init, sun_init)
-        plt.show()
-
-
-class ImageInteractive(NoInamge):
 
     def __init__(self, rfrac_init=0.5, sun_init=np.pi/2):
         super().__init__(rfrac_init, sun_init)
         self.init_interactivity(rfrac_init, sun_init)
-        plt.show()
 
     def init_interactivity(self, rfrac_init, sun_init):
         self.fig.subplots_adjust(left=0.25, bottom=0.25)
@@ -354,22 +358,34 @@ class ImageInteractive(NoInamge):
         self.sliders['rfrac'].reset()
 
 
-class AnimatedImage(NoInamge):
+class AnimatedImage(Image):
 
-    def __init__(self, rfrac_init=0.5, sun_init=-np.pi):
+    def __init__(self, rfrac_init=0.5, sun_init=-np.pi, runtime_sec=5, fps=30):
         super().__init__(rfrac_init, sun_init)
 
-        self.ani = animation.FuncAnimation(self.fig, self.animate, frames=150)
+        self.fig.tight_layout()
-        plt.show()
 
+        self.fps = fps
+        self.frame_num = self.fps * runtime_sec
+        self.ani = animation.FuncAnimation(self.fig, self.animate, frames=self.frame_num)
 
     def animate(self, frame_i):
-        phi = (frame_i*2*np.pi/150)-np.pi
+        phi = (frame_i*2*np.pi/self.frame_num)-np.pi
         self.update_sun(phi, 0)
         return self.lines.values()
 
+    def save(self, path="./torus.mp4"):
+        if path.endswith('.gif'):
+            writer = animation.PillowWriter(fps=self.fps)
+        else:
+            writer = animation.FFMpegWriter(fps=self.fps)
+        self.ani.save(path, writer)
+
 
 if __name__ == '__main__':
-    # ImageStatic()
+    img = Image()
-    # ImageInteractive()
+    # img = InteractiveImage()
-    AnimatedImage()
+    # img = AnimatedImage(runtime_sec=3, fps=10)
+
+    # plt.show()
+    img.save()