optional figure saving

1 files changed, 25 insertions, 1 deletions

diff --git a/mkcuts_dvpT.py b/mkcuts_dvpT.py
index cc9ed85..c7270c7 100644
--- a/mkcuts_dvpT.py
+++ b/mkcuts_dvpT.py
@@ -29,8 +29,15 @@ def main():
     parser.add_argument("-v", action="store_true", help='Show Velocity Cuts')
     parser.add_argument("-p", action="store_true", help='Show Pressure Cuts')
     parser.add_argument("-T", action="store_true", help='Show Temperature Cuts')
+    parser.add_argument("--save", action="store_true", help='Save Figures')
+    parser.add_argument("--show", action="store_true", help='Show Figures')
     args = parser.parse_args()
 
+    # Sanity Check
+    if not args.save and not args.show:
+        print "At least one of --save or --show is required."
+        sys.exit(0)
+
     # Give Feedback
     print "Creating Cuts for Output %i." % args.iout
     print ""
@@ -195,7 +202,24 @@ def main():
         plt.xlabel('Y')
         plt.ylabel('Z')
 
-    plt.show()
+    # Save Figures
+    if args.save:
+        if args.d:
+            f1.savefig('cut_dZ_%05d.png' % args.iout)
+            f2.savefig('cut_dX_%05d.png' % args.iout)
+        if args.v:
+            f3.savefig('cut_vZ_%05d.png' % args.iout)
+            f4.savefig('cut_vX_%05d.png' % args.iout)
+        if args.p:
+            f5.savefig('cut_pZ_%05d.png' % args.iout)
+            f6.savefig('cut_pX_%05d.png' % args.iout)
+        if args.T:
+            f7.savefig('cut_TZ_%05d.png' % args.iout)
+            f8.savefig('cut_TX_%05d.png' % args.iout)
+
+    # Show Figures?
+    if args.show:
+        plt.show()
 
 """
 Jump into Main().