Position keyword now adds position variables in PICMI ParticleDiagnos…

…tic (#4721)

* Position keyword in ParticleDiags now adds pos

Prior to recently, the `'position'` string in the input of the PICMI
`ParticleDiagnostic` did not add the positions because they were written out
automatically.
That has been changed and so we are adding them to the list of variables now.

Also added the class `LabFrameParticleDiagnostic` to the docs.

* Fix dimension-specific cases and examples

- in picmi.py, the dimensions were not checked and so it would try to add
  variables that do not exist for a species in a lower dimension
- in examples, some coordinates were named wrongly for lower dimensions
  than 2 but it did not matter before since positions were added automatically

Docs/source/usage/python.rst

@@ -94,6 +94,8 @@ Diagnostics
 
 Lab-frame diagnostics diagnostics are used when running boosted-frame simulations.
 
+.. autoclass:: pywarpx.picmi.LabFrameParticleDiagnostic
+
 .. autoclass:: pywarpx.picmi.LabFrameFieldDiagnostic
 
 Particles

Examples/Tests/ionization/PICMI_inputs_2d.py

@@ -88,7 +88,7 @@
     name = 'diag1',
     period = 10000,
     species = [electrons, ions],
-    data_list = ['ux', 'uy', 'uz', 'x', 'y', 'weighting'],
+    data_list = ['ux', 'uy', 'uz', 'x', 'z', 'weighting'],
     write_dir = '.',
     warpx_file_prefix = 'Python_ionization_plt')
 field_diag = picmi.FieldDiagnostic(

Examples/Tests/ohm_solver_ion_Landau_damping/PICMI_inputs.py

@@ -222,7 +222,7 @@ def setup_run(self):
                 period=100,
                 write_dir='.',
                 species=[self.ions],
-                data_list = ['ux', 'uy', 'uz', 'x', 'y', 'weighting'],
+                data_list = ['ux', 'uy', 'uz', 'x', 'z', 'weighting'],
                 warpx_file_prefix=f'Python_ohms_law_solver_landau_damping_{self.dim}d_plt',
             )
             simulation.add_diagnostic(particle_diag)

Examples/Tests/ohm_solver_ion_beam_instability/PICMI_inputs.py

@@ -259,7 +259,7 @@ def setup_run(self):
                 name='diag1',
                 period=1250,
                 species=[self.ions, self.beam_ions],
-                data_list = ['ux', 'uy', 'uz', 'x', 'weighting'],
+                data_list = ['ux', 'uy', 'uz', 'z', 'weighting'],
                 write_dir='.',
                 warpx_file_prefix='Python_ohms_law_solver_ion_beam_1d_plt',
             )

Examples/Tests/ohm_solver_magnetic_reconnection/PICMI_inputs.py

@@ -255,7 +255,7 @@ def setup_run(self):
                 period=self.total_steps,
                 write_dir='.',
                 species=[self.ions],
-                data_list=['ux', 'uy', 'uz', 'x', 'y', 'weighting'],
+                data_list=['ux', 'uy', 'uz', 'x', 'z', 'weighting'],
                 warpx_file_prefix='Python_ohms_law_solver_magnetic_reconnection_2d_plt',
                 # warpx_format='openpmd',
                 # warpx_openpmd_backend='h5',

Python/pywarpx/picmi.py

@@ -2442,8 +2442,13 @@ def diagnostic_initialize_inputs(self):
         if self.data_list is not None:
             for dataname in self.data_list:
                 if dataname == 'position':
-                    # --- The positions are alway written out anyway
-                    pass
+                    if pywarpx.geometry.dims != '1':  # because then it's WARPX_DIM_1D_Z
+                        variables.add('x')
+                    if pywarpx.geometry.dims == '3':
+                        variables.add('y')
+                    variables.add('z')
+                    if pywarpx.geometry.dims == 'RZ':
+                        variables.add('theta')
                 elif dataname == 'momentum':
                     variables.add('ux')
                     variables.add('uy')
@@ -2457,8 +2462,24 @@ def diagnostic_initialize_inputs(self):
                     variables.add('Bx')
                     variables.add('By')
                     variables.add('Bz')
-                elif dataname in ['ux', 'uy', 'uz', 'Ex', 'Ey', 'Ez', 'Bx', 'By', 'Bz']:
-                    variables.add(dataname)
+                elif dataname in ['x', 'y', 'z', 'theta', 'ux', 'uy', 'uz', 'Ex', 'Ey', 'Ez', 'Bx', 'By', 'Bz', 'Er', 'Et', 'Br', 'Bt']:
+                    if pywarpx.geometry.dims == '1' and (dataname == 'x' or dataname == 'y'):
+                        raise RuntimeError(
+                            f"The attribute {dataname} is not available in mode WARPX_DIM_1D_Z"
+                            f"chosen by dim={pywarpx.geometry.dims} in pywarpx."
+                        )
+                    elif pywarpx.geometry.dims != '3' and dataname == 'y':
+                        raise RuntimeError(
+                            f"The attribute {dataname} is not available outside of mode WARPX_DIM_3D"
+                            f"The chosen value was dim={pywarpx.geometry.dims} in pywarpx."
+                        )
+                    elif pywarpx.geometry.dims != 'RZ' and dataname == 'theta':
+                        raise RuntimeError(
+                            f"The attribute {dataname} is not available outside of mode WARPX_DIM_RZ."
+                            f"The chosen value was dim={pywarpx.geometry.dims} in pywarpx."
+                        )
+                    else:
+                        variables.add(dataname)
 
             # --- Convert the set to a sorted list so that the order
             # --- is the same on all processors.
@@ -2666,25 +2687,46 @@ def diagnostic_initialize_inputs(self):
         # --- Use a set to ensure that fields don't get repeated.
         variables = set()
 
-        if self.data_list is not None:
-            for dataname in self.data_list:
-                if dataname == 'position':
-                    # --- The positions are alway written out anyway
-                    pass
-                elif dataname == 'momentum':
-                    variables.add('ux')
-                    variables.add('uy')
-                    variables.add('uz')
-                elif dataname == 'weighting':
-                    variables.add('w')
-                elif dataname == 'fields':
-                    variables.add('Ex')
-                    variables.add('Ey')
-                    variables.add('Ez')
-                    variables.add('Bx')
-                    variables.add('By')
-                    variables.add('Bz')
-                elif dataname in ['ux', 'uy', 'uz', 'Ex', 'Ey', 'Ez', 'Bx', 'By', 'Bz', 'Er', 'Et', 'Br', 'Bt']:
+
+        for dataname in self.data_list:
+            if dataname == 'position':
+                if pywarpx.geometry.dims != '1':  # because then it's WARPX_DIM_1D_Z
+                    variables.add('x')
+                if pywarpx.geometry.dims == '3':
+                    variables.add('y')
+                variables.add('z')
+                if pywarpx.geometry.dims == 'RZ':
+                    variables.add('theta')
+            elif dataname == 'momentum':
+                variables.add('ux')
+                variables.add('uy')
+                variables.add('uz')
+            elif dataname == 'weighting':
+                variables.add('w')
+            elif dataname == 'fields':
+                variables.add('Ex')
+                variables.add('Ey')
+                variables.add('Ez')
+                variables.add('Bx')
+                variables.add('By')
+                variables.add('Bz')
+            elif dataname in ['x', 'y', 'z', 'theta', 'ux', 'uy', 'uz', 'Ex', 'Ey', 'Ez', 'Bx', 'By', 'Bz', 'Er', 'Et', 'Br', 'Bt']:
+                if pywarpx.geometry.dims == '1' and (dataname == 'x' or dataname == 'y'):
+                    raise RuntimeError(
+                        f"The attribute {dataname} is not available in mode WARPX_DIM_1D_Z"
+                        f"chosen by dim={pywarpx.geometry.dims} in pywarpx."
+                    )
+                elif pywarpx.geometry.dims != '3' and dataname == 'y':
+                    raise RuntimeError(
+                        f"The attribute {dataname} is not available outside of mode WARPX_DIM_3D"
+                        f"The chosen value was dim={pywarpx.geometry.dims} in pywarpx."
+                    )
+                elif pywarpx.geometry.dims != 'RZ' and dataname == 'theta':
+                    raise RuntimeError(
+                        f"The attribute {dataname} is not available outside of mode WARPX_DIM_RZ."
+                        f"The chosen value was dim={pywarpx.geometry.dims} in pywarpx."
+                    )
+                else:
                     variables.add(dataname)
 
             # --- Convert the set to a sorted list so that the order