Implement base openPMD diagnostics

examples/track_plasma_fluid.py

@@ -32,13 +32,14 @@
 
 # Create plasma stage.
 plasma = PlasmaStage(
-    1e-2, 1e23, laser=laser, wakefield_model='cold_fluid_1d', n_out=20,
+    1e-2, 1e23, laser=laser, wakefield_model='cold_fluid_1d', n_out=50,
     laser_evolution=True, laser_z_foc=0, beam_wakefields=True,
     r_max=70e-6,  xi_min=40e-6, xi_max=120e-6, n_r=70, n_xi=200)
 
 
 # Do tracking.
-bunch_list = plasma.track(bunch, out_initial=True)
+opmd_diag = False  # Set to True to active openPMD output.
+bunch_list = plasma.track(bunch, out_initial=True, opmd_diag=opmd_diag)
 
 
 # Analyze bunch evolution.

examples/track_plasma_qs2d.py

@@ -41,7 +41,8 @@
 
 
 # Do tracking.
-bunch_list = plasma.track(bunch, out_initial=True)
+opmd_diag = False  # Set to True to active openPMD output.
+bunch_list = plasma.track(bunch, out_initial=True, opmd_diag=opmd_diag)
 
 
 # Analyze bunch evolution.

requirements.txt

@@ -1,4 +1,5 @@
 numpy
 scipy
 numba
-aptools
+aptools
+openpmd-api

wake_t/__version__.py

@@ -1 +1 @@
-__version__ = "0.4.2"
+__version__ = "0.4.3"

wake_t/beamline_elements.py

@@ -17,6 +17,7 @@
 from wake_t.utilities.bunch_manipulation import (
     convert_to_ocelot_matrix, convert_from_ocelot_matrix, rotation_matrix_xz)
 from wake_t.csr import get_csr_calculator
+from wake_t.diagnostics import OpenPMDDiagnostics
 
 
 class Beamline():
@@ -26,7 +27,7 @@ class Beamline():
     def __init__(self, elements):
         self.elements = elements
 
-    def track(self, bunch, out_initial=True):
+    def track(self, bunch, out_initial=True, opmd_diag=False, diag_dir=None):
         """
         Track bunch through beamline.
 
@@ -40,6 +41,19 @@ def track(self, bunch, out_initial=True):
             output bunch list. This applies only at the beginning and not for
             every beamline element.
 
+        opmd_diag : bool or OpenPMDDiagnostics
+            Determines whether to write simulation diagnostics to disk (i.e.
+            particle distributions and fields). The output is written to
+            HDF5 files following the openPMD standard. The number of outputs
+            per beamline element is determined by its `n_out` value. It is also
+            possible to provide an already existing OpenPMDDiagnostics
+            instance instead of a boolean value.
+
+        diag_dir : str
+            Directory into which the openPMD output will be written. By default
+            this is a 'diags' folder in the current directory. Only needed if
+            `opmd_diag=True`.
+
         Returns:
         --------
         A list of size 'n_out' containing the bunch distribution at each step.
@@ -48,8 +62,11 @@ def track(self, bunch, out_initial=True):
         bunch_list = []
         if out_initial:
             bunch_list.append(copy(bunch))
+        if type(opmd_diag) is not OpenPMDDiagnostics and opmd_diag:
+            opmd_diag = OpenPMDDiagnostics(write_dir=diag_dir)
         for element in self.elements:
-            bunch_list.extend(element.track(bunch, out_initial=False))
+            bunch_list.extend(
+                element.track(bunch, out_initial=False, opmd_diag=opmd_diag))
         return bunch_list
 
 
@@ -252,7 +269,8 @@ def __init__(self, length, n_p, laser=None, tracking_mode='numerical',
         self.n_out = n_out
         self.model_params = model_params
 
-    def track(self, bunch, parallel=False, n_proc=None, out_initial=False):
+    def track(self, bunch, parallel=False, n_proc=None, out_initial=False,
+              opmd_diag=False, diag_dir=None):
         """
         Track bunch through plasma stage.
 
@@ -272,6 +290,18 @@ def track(self, bunch, parallel=False, n_proc=None, out_initial=False):
             Determines whether the initial bunch should be included in the
             output bunch list.
 
+        opmd_diag : bool or OpenPMDDiagnostics
+            Determines whether to write simulation diagnostics to disk (i.e.
+            particle distributions and fields). The output is written to
+            HDF5 files following the openPMD standard. The number of outputs
+            the `n_out` value. It is also possible to provide an already
+            existing OpenPMDDiagnostics instance instead of a boolean value.
+
+        diag_dir : str
+            Directory into which the openPMD output will be written. By default
+            this is a 'diags' folder in the current directory. Only needed if
+            `opmd_diag=True`.
+
         Returns:
         --------
         A list of size 'n_out' containing the bunch distribution at each step.
@@ -282,12 +312,17 @@ def track(self, bunch, parallel=False, n_proc=None, out_initial=False):
         print('-'*len('Plasma stage'))
         if out_initial:
             initial_bunch = copy(bunch)
+        if type(opmd_diag) is not OpenPMDDiagnostics and opmd_diag:
+            opmd_diag = OpenPMDDiagnostics(write_dir=diag_dir)
         if self.tracking_mode == 'numerical':
-            bunch_list = self._track_numerically(bunch, parallel, n_proc)
+            bunch_list = self._track_numerically(
+                bunch, parallel, n_proc, opmd_diag)
         elif self.tracking_mode == 'analytical':
             bunch_list = self._track_analytically(bunch, parallel, n_proc)
         if out_initial:
             bunch_list.insert(0, initial_bunch)
+        if opmd_diag is not False:
+            opmd_diag.increase_z_pos(self.length)
         return bunch_list
 
     def _get_wakefield(self, wakefield_model, model_params):
@@ -315,7 +350,7 @@ def _get_wakefield(self, wakefield_model, model_params):
     def _gamma(self, px, py, pz):
         return np.sqrt(1 + px**2 + py**2 + pz**2)
 
-    def _track_numerically(self, bunch, parallel, n_proc):
+    def _track_numerically(self, bunch, parallel, n_proc, opmd_diag):
         # Get 6D matrix
         mat = bunch.get_6D_matrix_with_charge()
         # Plasma length in time
@@ -365,6 +400,9 @@ def _track_numerically(self, bunch, parallel, n_proc):
                         ParticleBunch(bunch.q, x, y, xi, px, py, pz,
                                       prop_distance=new_prop_dist)
                     )
+                    if opmd_diag is not False:
+                        opmd_diag.write_diagnostics(
+                            s*t_step, t_step, [bunch_list[-1]], self.wakefield)
             finally:
                 process_pool.close()
                 process_pool.join()
@@ -386,6 +424,9 @@ def _track_numerically(self, bunch, parallel, n_proc):
                     ParticleBunch(bunch.q, x, y, xi, px, py, pz,
                                   prop_distance=new_prop_dist)
                 )
+                if opmd_diag is not False:
+                    opmd_diag.write_diagnostics(
+                        s*t_step, t_step, [bunch_list[-1]], self.wakefield)
         # print computing time
         end = time.time()
         print("Done ({:1.3f} seconds).".format(end-start))
@@ -693,7 +734,8 @@ def __init__(self, length, plasma_dens_top, plasma_dens_down=None,
         self.n_out = n_out
         self.wakefield = self._get_wakefield(wakefield_model, model_params)
 
-    def track(self, bunch, parallel=False, n_proc=None, out_initial=False):
+    def track(self, bunch, parallel=False, n_proc=None, out_initial=False,
+              opmd_diag=False, diag_dir=None):
         """
         Track bunch through plasma ramp.
 
@@ -713,6 +755,18 @@ def track(self, bunch, parallel=False, n_proc=None, out_initial=False):
             Determines whether the initial bunch should be included in the
             output bunch list.
 
+        opmd_diag : bool or OpenPMDDiagnostics
+            Determines whether to write simulation diagnostics to disk (i.e.
+            particle distributions and fields). The output is written to
+            HDF5 files following the openPMD standard. The number of outputs
+            the `n_out` value. It is also possible to provide an already
+            existing OpenPMDDiagnostics instance instead of a boolean value.
+
+        diag_dir : str
+            Directory into which the openPMD output will be written. By default
+            this is a 'diags' folder in the current directory. Only needed if
+            `opmd_diag=True`.
+
         Returns:
         --------
         A list of size 'n_out' containing the bunch distribution at each step.
@@ -735,6 +789,8 @@ def track(self, bunch, parallel=False, n_proc=None, out_initial=False):
         bunch_list = list()
         if out_initial:
             bunch_list.append(copy(bunch))
+        if type(opmd_diag) is not OpenPMDDiagnostics and opmd_diag:
+            opmd_diag = OpenPMDDiagnostics(write_dir=diag_dir)
         start = time.time()
 
         if parallel:
@@ -766,6 +822,9 @@ def track(self, bunch, parallel=False, n_proc=None, out_initial=False):
                         ParticleBunch(bunch.q, x, y, xi, px, py, pz,
                                       prop_distance=new_prop_dist)
                     )
+                    if opmd_diag is not False:
+                        opmd_diag.write_diagnostics(
+                            s*t_step, t_step, [bunch_list[-1]], self.wakefield)
             finally:
                 process_pool.close()
                 process_pool.join()
@@ -784,6 +843,9 @@ def track(self, bunch, parallel=False, n_proc=None, out_initial=False):
                     ParticleBunch(bunch.q, x, y, xi, px, py, pz,
                                   prop_distance=new_prop_dist)
                 )
+                if opmd_diag is not False:
+                    opmd_diag.write_diagnostics(
+                        s*t_step, t_step, [bunch_list[-1]], self.wakefield)
         end = time.time()
         print("Done ({:1.3f} seconds).".format(end-start))
         print('-'*80)
@@ -792,6 +854,8 @@ def track(self, bunch, parallel=False, n_proc=None, out_initial=False):
         bunch.set_phase_space(last_bunch.x, last_bunch.y, last_bunch.xi,
                               last_bunch.px, last_bunch.py, last_bunch.pz)
         bunch.increase_prop_distance(self.length)
+        if opmd_diag is not False:
+            opmd_diag.increase_z_pos(self.length)
 
         return bunch_list
 
@@ -1014,7 +1078,8 @@ def _get_wakefield(self, relativistic, wakefields, wakefield_model,
     def calculate_density(self, z):
         return self.n_p
 
-    def track(self, bunch, parallel=False, n_proc=None, out_initial=False):
+    def track(self, bunch, parallel=False, n_proc=None, out_initial=False,
+              opmd_diag=False, diag_dir=None):
         """
         Track bunch through plasma lens.
 
@@ -1034,6 +1099,18 @@ def track(self, bunch, parallel=False, n_proc=None, out_initial=False):
             Determines whether the initial bunch should be included in the
             output bunch list.
 
+        opmd_diag : bool or OpenPMDDiagnostics
+            Determines whether to write simulation diagnostics to disk (i.e.
+            particle distributions and fields). The output is written to
+            HDF5 files following the openPMD standard. The number of outputs
+            the `n_out` value. It is also possible to provide an already
+            existing OpenPMDDiagnostics instance instead of a boolean value.
+
+        diag_dir : str
+            Directory into which the openPMD output will be written. By default
+            this is a 'diags' folder in the current directory. Only needed if
+            `opmd_diag=True`.
+
         Returns:
         --------
         A list of size 'n_out' containing the bunch distribution at each step.
@@ -1057,6 +1134,8 @@ def track(self, bunch, parallel=False, n_proc=None, out_initial=False):
         bunch_list = list()
         if out_initial:
             bunch_list.append(copy(bunch))
+        if type(opmd_diag) is not OpenPMDDiagnostics and opmd_diag:
+            opmd_diag = OpenPMDDiagnostics(write_dir=diag_dir)
         start = time.time()
         if parallel:
             if n_proc is None:
@@ -1086,6 +1165,9 @@ def track(self, bunch, parallel=False, n_proc=None, out_initial=False):
                     bunch_list.append(
                         ParticleBunch(bunch.q, x, y, xi, px, py, pz,
                                       prop_distance=new_prop_dist))
+                    if opmd_diag is not False:
+                        opmd_diag.write_diagnostics(
+                            s*t_step, t_step, [bunch_list[-1]], self.field)
             finally:
                 process_pool.close()
                 process_pool.join()
@@ -1105,6 +1187,9 @@ def track(self, bunch, parallel=False, n_proc=None, out_initial=False):
                     ParticleBunch(bunch.q, x, y, xi, px, py, pz,
                                   prop_distance=new_prop_dist)
                 )
+                if opmd_diag is not False:
+                    opmd_diag.write_diagnostics(
+                        s*t_step, t_step, [bunch_list[-1]], self.field)
         end = time.time()
         print("Done ({:1.3f} seconds).".format(end-start))
         print('-'*80)
@@ -1113,6 +1198,8 @@ def track(self, bunch, parallel=False, n_proc=None, out_initial=False):
         bunch.set_phase_space(last_bunch.x, last_bunch.y, last_bunch.xi,
                               last_bunch.px, last_bunch.py, last_bunch.pz)
         bunch.increase_prop_distance(self.length)
+        if opmd_diag is not False:
+            opmd_diag.increase_z_pos(self.length)
         return bunch_list
 
     def _get_optimized_dt(self, beam, WF):
@@ -1144,7 +1231,8 @@ def __init__(self, length=0, theta=0, k1=0, k2=0, gamma_ref=None,
         self.csr_calculator = get_csr_calculator()
         self.element_name = ''
 
-    def track(self, bunch, backtrack=False, out_initial=False):
+    def track(self, bunch, backtrack=False, out_initial=False, opmd_diag=False,
+              diag_dir=None):
         # Convert bunch to ocelot units and reference frame
         bunch_mat, g_avg = self._get_beam_matrix_for_tracking(bunch)
         if self.gamma_ref is None:
@@ -1158,6 +1246,10 @@ def track(self, bunch, backtrack=False, out_initial=False):
         # Determine track and output steps
         l_step, track_steps, output_steps = self._determine_steps()
 
+        # Create diagnostics if needed
+        if type(opmd_diag) is not OpenPMDDiagnostics and opmd_diag:
+            opmd_diag = OpenPMDDiagnostics(write_dir=diag_dir)
+
         # Print output header
         print('')
         print(self.element_name.capitalize())
@@ -1192,10 +1284,17 @@ def track(self, bunch, backtrack=False, out_initial=False):
                 new_bunch = self._create_new_bunch(
                     bunch, new_bunch_mat, l_curr)
                 output_bunch_list.append(new_bunch)
+                if opmd_diag is not False:
+                    opmd_diag.write_diagnostics(
+                        l_curr/ct.c, l_step/ct.c, [output_bunch_list[-1]])
 
         # Update bunch data
         self._update_input_bunch(bunch, bunch_mat, output_bunch_list)
 
+        # Add element length to diagnostics position
+        if opmd_diag is not False:
+            opmd_diag.increase_z_pos(self.length)
+
         # Finalize
         tracking_time = time.time() - start_time
         print('Done ({} s).'.format(tracking_time))