Adding an if statement for the last step diagnostic (#3436)

* modified

* Include max_step and stop_time as input parameters

* Add comments

* update

* update

* Fix CI test: apply suggestion from @ax3l

* Do not force consistency between `max_step` and `stop_time`

Co-authored-by: Dave Grote <grote1@llnl.gov>
Co-authored-by: Edoardo Zoni <ezoni@lbl.gov>

2 files changed, 11 insertions, 14 deletions

diff --git a/Python/pywarpx/WarpX.py b/Python/pywarpx/WarpX.py
index 517c74d42..76f66c6d2 100644
--- a/Python/pywarpx/WarpX.py
+++ b/Python/pywarpx/WarpX.py
@@ -30,8 +30,13 @@ class WarpX(Bucket):
     A Python wrapper for the WarpX C++ class
     """
 
-    def create_argv_list(self):
+    def create_argv_list(self, **kw):
         argv = []
+
+        for k, v in kw.items():
+            if v is not None:
+                argv.append(f'{k} = {v}')
+
         argv += warpx.attrlist()
         argv += my_constants.attrlist()
         argv += amr.attrlist()
@@ -86,8 +91,8 @@ class WarpX(Bucket):
 
         return argv
 
-    def init(self, mpi_comm=None):
-        argv = ['warpx'] + self.create_argv_list()
+    def init(self, mpi_comm=None, **kw):
+        argv = ['warpx'] + self.create_argv_list(**kw)
         libwarpx.initialize(argv, mpi_comm=mpi_comm)
 
     def evolve(self, nsteps=-1):
@@ -103,10 +108,7 @@ class WarpX(Bucket):
         return libwarpx.libwarpx_so.warpx_getProbHi(direction)
 
     def write_inputs(self, filename='inputs', **kw):
-        argv = self.create_argv_list()
-
-        for k, v in kw.items():
-            argv.append(f'{k} = {v}')
+        argv = self.create_argv_list(**kw)
 
         # Sort the argv list to make it more human readable
         argv.sort()
diff --git a/Python/pywarpx/picmi.py b/Python/pywarpx/picmi.py
index 0d6eebf63..9399fd926 100644
--- a/Python/pywarpx/picmi.py
+++ b/Python/pywarpx/picmi.py
@@ -1586,16 +1586,11 @@ class Simulation(picmistandard.PICMI_Simulation):
             return
 
         self.warpx_initialized = True
-        pywarpx.warpx.init(mpi_comm)
+        pywarpx.warpx.init(mpi_comm, max_step=self.max_steps, stop_time=self.max_time)
 
     def write_input_file(self, file_name='inputs'):
         self.initialize_inputs()
-        kw = {}
-        if self.max_steps is not None:
-            kw['max_step'] = self.max_steps
-        if self.max_time is not None:
-            kw['stop_time'] = self.max_time
-        pywarpx.warpx.write_inputs(file_name, **kw)
+        pywarpx.warpx.write_inputs(file_name, max_step=self.max_steps, stop_time=self.max_time)
 
     def step(self, nsteps=None, mpi_comm=None):
         self.initialize_inputs()