Order-independent `rho_<species>` with RZ and openPMD (#3338)

* order-independent rho_<species>

* add test to check rho_<species> is order-agnostic

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* make analysis naming more consistent

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

1 files changed, 24 insertions, 1 deletions

diff --git a/Examples/Tests/openpmd_rz/analysis_openpmd_rz.py b/Examples/Tests/openpmd_rz/analysis_openpmd_rz.py
index 4cd48a2c6..247c4ac61 100755
--- a/Examples/Tests/openpmd_rz/analysis_openpmd_rz.py
+++ b/Examples/Tests/openpmd_rz/analysis_openpmd_rz.py
@@ -1,5 +1,6 @@
 #!/usr/bin/env python3
 
+import numpy as np
 import openpmd_api as io
 
 series = io.Series("LaserAccelerationRZ_opmd_plt/openpmd_%T.h5", io.Access.read_only)
@@ -8,7 +9,7 @@ assert len(series.iterations) == 3, 'improper number of iterations stored'
 
 ii = series.iterations[20]
 
-assert len(ii.meshes) == 7, 'improper number of meshes'
+assert len(ii.meshes) == 8, 'improper number of meshes'
 
 # select j_t
 jt = ii.meshes['j']['t']
@@ -22,3 +23,25 @@ assert Nz == 512, 'Wrong number of z points stored or possible incorrect orderin
 
 assert ii.meshes['part_per_grid'][io.Mesh_Record_Component.SCALAR].shape == [512,64], 'problem with part_per_grid'
 assert ii.meshes['rho_electrons'][io.Mesh_Record_Component.SCALAR].shape == [3, 512, 64], 'problem with rho_electrons'
+
+
+### test that openpmd+RZ
+### 1. creates rho per species correctly
+### 2. orders these appropriately
+rhoe_mesh = ii.meshes['rho_electrons']
+rhob_mesh = ii.meshes['rho_beam']
+dz, dr = rhoe_mesh.grid_spacing
+zmin, rmin = rhoe_mesh.grid_global_offset
+
+rhoe = rhoe_mesh[io.Mesh_Record_Component.SCALAR][:]
+rhob = rhob_mesh[io.Mesh_Record_Component.SCALAR][:]
+series.flush()
+nm, nz, nr = rhoe.shape
+zlist = zmin + dz * np.arange(nz)
+rhoe0 = rhoe[0] # 0 mode
+rhob0 = rhob[0] # 0 mode
+
+electron_meanz = np.sum(np.dot(zlist, rhoe0))/ np.sum(rhoe0)
+beam_meanz = np.sum(np.dot(zlist, rhob0))/ np.sum(rhob0)
+
+assert ((electron_meanz > 0) and (beam_meanz < 0)), 'problem with openPMD+RZ.  Maybe openPMD+RZ mixed up the order of rho_<species> diagnostics?'