add argparse

torus.py

@@ -32,11 +32,12 @@ class TorusWorld:
         c0 = k3**2 - 4*(rmaj**2)*(rmin**2 - o[2]**2)
         return [c0, c1, c2, c3, c4]
 
-    def __init__(self, rfrac):
+    def __init__(self, rfrac, sun=np.pi/2):
         self.r_maj = 1
         self.r_min = rfrac
-        self.sun = np.array([self.r_maj, 0, self.r_maj])
+        self.sun = None
-        self.sun_r = np.array([np.pi/2, 0])
+        self.sun_r = None
+        self.put_sun(sun, 0)
         self.surface_map = None
 
     def update(self, rfrac=None):
@@ -44,7 +45,8 @@ class TorusWorld:
         self.r_min = rfrac
         self.put_sun(*self.sun_r)
 
-    def put_sun(self, phi, theta):
+    def put_sun(self, phi_in, theta_in):
+        phi, theta = (np.array([phi_in, theta_in]) + np.pi) % (2 * np.pi) - np.pi
         self.sun_r = np.array([phi, theta])
         self.sun = np.array([
             (self.r_maj - self.r_maj*np.cos(phi))*np.cos(theta),
@@ -98,7 +100,7 @@ class Image():
         self.contour_kwargs = dict(cmap=plt.colormaps['hot'], vmin=0, vmax=1)
         self.levels = 25  # [-1, 0, 1]
 
-        self.torus = TorusWorld(rfrac_init)
+        self.torus = TorusWorld(rfrac_init, sun_init)
         self.illumination = None
         self.res = 200
         self.rectangular_map = True
@@ -367,10 +369,11 @@ class AnimatedImage(Image):
 
         self.fps = fps
         self.frame_num = self.fps * runtime_sec
+        self.sun_init = sun_init
         self.ani = animation.FuncAnimation(self.fig, self.animate, frames=self.frame_num)
 
     def animate(self, frame_i):
-        phi = (frame_i*2*np.pi/self.frame_num)-np.pi
+        phi = (frame_i*2*np.pi/self.frame_num) + self.sun_init
         self.update_sun(phi, 0)
         return self.lines.values()
 
@@ -383,9 +386,79 @@ class AnimatedImage(Image):
 
 
 if __name__ == '__main__':
-    img = Image()
+    import argparse
-    # img = InteractiveImage()
+    import pathlib
-    # img = AnimatedImage(runtime_sec=3, fps=10)
 
-    # plt.show()
+    def parse_args():
-    img.save()
+
+        def create_static(args):
+            img = Image(args.r_minor, args.sun)
+
+            if args.output:
+                img.save(args.output)
+            else:
+                plt.show()
+
+        def create_interactive(args):
+            InteractiveImage(args.r_minor, args.sun)
+            plt.show()
+
+        def create_animation(args):
+            img = AnimatedImage(args.r_minor, args.sun, args.runtime, args.fps)
+
+            if args.output:
+                img.save(args.output)
+            else:
+                plt.show()
+
+        class Range(object):
+            def __init__(self, start, end):
+                self.start = start
+                self.end = end
+
+            def __eq__(self, other):
+                return self.start <= other <= self.end
+
+            def __contains__(self, item):
+                return self.__eq__(item)
+
+            def __iter__(self):
+                yield self
+
+            def __str__(self):
+                return '[{0},{1}]'.format(self.start, self.end)
+
+        p = argparse.ArgumentParser()
+        sp = p.add_subparsers()
+        p_static = sp.add_parser('static', aliases=['s'], help='create a static image')
+        p_interactive = sp.add_parser('interactive', aliases=['i'], help='create an interactive image')
+        p_animated = sp.add_parser('animated', aliases=['a'], help='create a animated image')
+
+        p_static.add_argument('r_minor', help='Ratio of minor radius to major radius',
+                              nargs='?', type=float, choices=Range(0., 1.), metavar='RMIN', default=0.5)
+        p_static.add_argument('sun', help="Position of the sun as an angle measured from the torus' origin",
+                              nargs='?', type=float, metavar='PHI', default=np.pi/2)
+        p_static.add_argument('--output', '-o', help='output to a file instead of displaying directly',
+                              type=pathlib.Path, metavar='FILE')
+        p_static.set_defaults(func=create_static)
+
+        p_interactive.add_argument('r_minor', help='Ratio of minor radius to major radius',
+                                   nargs='?', type=float, choices=Range(0., 1.), metavar='RMIN', default=0.5)
+        p_interactive.add_argument('sun', help="Position of the sun as an angle measured from the torus' origin",
+                                   nargs='?', type=float, metavar='PHI', default=np.pi/2)
+        p_interactive.set_defaults(func=create_interactive)
+
+        p_animated.add_argument('r_minor', help='Ratio of minor radius to major radius',
+                                nargs='?', type=float, choices=Range(0., 1.), metavar='RMIN', default=0.5)
+        p_animated.add_argument('sun', help="Position of the sun as an angle measured from the torus' origin",
+                                nargs='?', type=float, metavar='PHI', default=np.pi/2)
+        p_animated.add_argument('--output', '-o', help='output to a file instead of displaying directly',
+                                metavar='FILE')
+        p_animated.add_argument('--fps', '-f', help='framerate of animation', metavar='FPS', type=int, default=20)
+        p_animated.add_argument('--runtime', '-r', help='runtime of animation', metavar='SEC', type=int, default=5)
+        p_animated.set_defaults(func=create_animation)
+
+        return (p.parse_args())
+
+    args = parse_args()
+    args.func(args)