make the python driver usable from the interpreter.

2 files changed, 32 insertions, 19 deletions

diff --git a/Exec/Langmuir/python/Langmuir.py b/Exec/Langmuir/python/Langmuir.py
index 672da52a4..6a6e4f2f6 100755
--- a/Exec/Langmuir/python/Langmuir.py
+++ b/Exec/Langmuir/python/Langmuir.py
@@ -3,23 +3,22 @@
 import matplotlib.pyplot as plt
 import warpx
 
-warpx.initialize()
+with warpx.script():
 
-# run for ten time steps
-warpx.evolve()
+    # run for ten time steps
+    warpx.evolve(10)
 
-x = warpx.get_particle_x(0)
-y = warpx.get_particle_y(0)
-plt.plot(x[0], y[0], '.')
-plt.savefig("particles.png")
+    x = warpx.get_particle_x(0)
+    y = warpx.get_particle_y(0)
+    plt.plot(x[0], y[0], '.')
+    plt.savefig("particles.png")
 
-# this returns a list of numpy arrays that hold the magnetic field 
-# data in the x-direction on each grid for level 0
-grid_data = warpx.get_mesh_magnetic_field(0, 0, False)
+    # this returns a list of numpy arrays that hold the magnetic field 
+    # data in the x-direction on each grid for level 0
+    grid_data = warpx.get_mesh_magnetic_field(0, 0, False)
 
-# plot a slice through the second grid 
-plt.clf()
-plt.pcolormesh(grid_data[1][9,:,:])
-plt.savefig("field.png")
+    # plot a slice through the second grid 
+    plt.clf()
+    plt.pcolormesh(grid_data[1][9,:,:])
+    plt.savefig("field.png")
 
-warpx.finalize()
diff --git a/Exec/Langmuir/python/warpx.py b/Exec/Langmuir/python/warpx.py
index d8b97f202..5662a4fd7 100644
--- a/Exec/Langmuir/python/warpx.py
+++ b/Exec/Langmuir/python/warpx.py
@@ -3,6 +3,7 @@ import ctypes
 from ctypes.util import find_library
 import numpy as np
 from numpy.ctypeslib import ndpointer
+from contextlib import contextmanager
 
 libwarpx = ctypes.CDLL("libwarpx.so")
 libc = ctypes.CDLL(find_library('c'))
@@ -78,9 +79,21 @@ def print_step(i):
 
     '''
     print("\nCalling a Python function from C++.  Step = ", i)
+
+
 c_print_step = CALLBACK_FUNC_1(print_step)
 
-def initialize():
+
+@contextmanager
+def script():
+    if len(sys.argv) < 2:
+        print("Usage: python <script_name> <inputs> <extra_args>")
+        sys.exit()
+    initialize(sys.argv[1:])
+    yield
+    finalize()
+
+def initialize(args):
     '''
     
     Initialize WarpX and AMReX. Must be called before 
@@ -88,10 +101,11 @@ def initialize():
     
     '''
     
-    # convert command line args to pass into amrex
-    argc = len(sys.argv)
+    # convert args to pass into amrex
+    cmd_line_args = ["main"] + args
+    argc = len(cmd_line_args)
     argv = (LP_c_char * (argc+1))()
-    for i, arg in enumerate(sys.argv):
+    for i, arg in enumerate(cmd_line_args):
         enc_arg = arg.encode('utf-8')
         argv[i] = ctypes.create_string_buffer(enc_arg)