Add by-cell averages of particle properties (#2892)

* Add ParticleReductionFunctor

* Adds a by-cell particle average diagnostic
* Averaged quantity is given by an input parser in (x,y,z,ux,uy,uz)
* Coordinates in these parsers are in the WarpX convention
* Users specify which species the averages are calculated for. The
averages are calculated for each given species individually.
* If there are no particles in a cell, the average will be zero.

* add documentation for ParticleReductionFunctor

* add test for cell-averaged particle diagnostics

* set benchmarks for praticle_fields_diags and particle_fields_diags_single_precision

* small changes to parameter docs and test script (addressing David's comments)

* Apply suggestions from code review

* remove unnecessary comments in python test script

* remove unused variable

* fixed inaccurate comments

* use static_cast<int> pattern

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

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

* add comment in ParticleReductionFunctor.H

* warpx.serialize_ics -> warpx.serialize_initial_conditions

* AMREX_SPACEDIM to WARPX_DIM_3D and WARPX_DIM_XZ

* fix compile errors from dimensionality conditionals

* fix compile errors

* Apply suggestions from code review

* clarify docs

* add include headers

* update docs

Co-authored-by: Edoardo Zoni <59625522+EZoni@users.noreply.github.com>
Co-authored-by: Axel Huebl <axel.huebl@plasma.ninja>

1 files changed, 156 insertions, 0 deletions

diff --git a/Examples/Tests/particle_fields_diags/analysis_particle_diags_impl.py b/Examples/Tests/particle_fields_diags/analysis_particle_diags_impl.py
new file mode 100755
index 000000000..e6dfc6a85
--- /dev/null
+++ b/Examples/Tests/particle_fields_diags/analysis_particle_diags_impl.py
@@ -0,0 +1,156 @@
+#!/usr/bin/env python3
+
+# Copyright 2019-2022 Luca Fedeli, Yinjian Zhao, Hannah Klion
+#
+# This file is part of WarpX.
+#
+# License: BSD-3-Clause-LBNL
+
+# This script tests the reduced particle diagnostics.
+# The setup is a uniform plasma with electrons, protons and photons.
+# Various particle and field quantities are written to file using the reduced diagnostics
+# and compared with the corresponding quantities computed from the data in the plotfiles.
+
+import os
+import sys
+
+import numpy as np
+from scipy.constants import c
+from scipy.constants import epsilon_0 as eps0
+from scipy.constants import m_e, m_p
+from scipy.constants import mu_0 as mu0
+import yt
+
+sys.path.insert(1, '../../../../warpx/Regression/Checksum/')
+import checksumAPI
+
+
+def do_analysis(single_precision = False):
+    fn = sys.argv[1]
+
+    ds = yt.load(fn)
+    ad = ds.all_data()
+    ad0 = ds.covering_grid(level=0, left_edge=ds.domain_left_edge, dims=ds.domain_dimensions)
+
+    #--------------------------------------------------------------------------------------------------
+    # Part 1: get results from plotfiles (label '_yt')
+    #--------------------------------------------------------------------------------------------------
+
+    # Quantities computed from plotfiles
+    values_yt = dict()
+
+    domain_size = ds.domain_right_edge.value - ds.domain_left_edge.value
+    dx = domain_size / ds.domain_dimensions
+
+    # Electrons
+    x = ad['electrons', 'particle_position_x'].to_ndarray()
+    y = ad['electrons', 'particle_position_y'].to_ndarray()
+    z = ad['electrons', 'particle_position_z'].to_ndarray()
+    uz = ad['electrons', 'particle_momentum_z'].to_ndarray() / m_e / c
+    w  = ad['electrons', 'particle_weight'].to_ndarray()
+
+    x_ind = ((x - ds.domain_left_edge[0].value) / dx[0]).astype(int)
+    y_ind = ((y - ds.domain_left_edge[1].value) / dx[1]).astype(int)
+    z_ind = ((z - ds.domain_left_edge[2].value) / dx[2]).astype(int)
+
+    zavg = np.zeros(ds.domain_dimensions)
+    uzavg = np.zeros(ds.domain_dimensions)
+    zuzavg = np.zeros(ds.domain_dimensions)
+    wavg = np.zeros(ds.domain_dimensions)
+
+    for i_p in range(len(x)):
+        zavg[x_ind[i_p],y_ind[i_p],z_ind[i_p]] += z[i_p] * w[i_p]
+        uzavg[x_ind[i_p],y_ind[i_p],z_ind[i_p]] += uz[i_p] * w[i_p]
+        zuzavg[x_ind[i_p],y_ind[i_p],z_ind[i_p]] += z[i_p] * uz[i_p] * w[i_p]
+        wavg[x_ind[i_p],y_ind[i_p],z_ind[i_p]] += w[i_p]
+
+    wavg_adj = np.where(wavg == 0, 1, wavg)
+    values_yt['electrons: zavg'] = zavg / wavg_adj
+    values_yt['electrons: uzavg'] = uzavg / wavg_adj
+    values_yt['electrons: zuzavg'] = zuzavg / wavg_adj
+
+    # protons
+    x = ad['protons', 'particle_position_x'].to_ndarray()
+    y = ad['protons', 'particle_position_y'].to_ndarray()
+    z = ad['protons', 'particle_position_z'].to_ndarray()
+    uz = ad['protons', 'particle_momentum_z'].to_ndarray() / m_p / c
+    w  = ad['protons', 'particle_weight'].to_ndarray()
+
+    x_ind = ((x - ds.domain_left_edge[0].value) / dx[0]).astype(int)
+    y_ind = ((y - ds.domain_left_edge[1].value) / dx[1]).astype(int)
+    z_ind = ((z - ds.domain_left_edge[2].value) / dx[2]).astype(int)
+
+    zavg = np.zeros(ds.domain_dimensions)
+    uzavg = np.zeros(ds.domain_dimensions)
+    zuzavg = np.zeros(ds.domain_dimensions)
+    wavg = np.zeros(ds.domain_dimensions)
+
+    for i_p in range(len(x)):
+        zavg[x_ind[i_p],y_ind[i_p],z_ind[i_p]] += z[i_p] * w[i_p]
+        uzavg[x_ind[i_p],y_ind[i_p],z_ind[i_p]] += uz[i_p] * w[i_p]
+        zuzavg[x_ind[i_p],y_ind[i_p],z_ind[i_p]] += z[i_p] * uz[i_p] * w[i_p]
+        wavg[x_ind[i_p],y_ind[i_p],z_ind[i_p]] += w[i_p]
+
+    wavg_adj = np.where(wavg == 0, 1, wavg)
+    values_yt['protons: zavg'] = zavg / wavg_adj
+    values_yt['protons: uzavg'] = uzavg / wavg_adj
+    values_yt['protons: zuzavg'] = zuzavg / wavg_adj
+
+    # Photons (momentum in units of m_e c)
+    x = ad['photons', 'particle_position_x'].to_ndarray()
+    y = ad['photons', 'particle_position_y'].to_ndarray()
+    z = ad['photons', 'particle_position_z'].to_ndarray()
+    uz = ad['photons', 'particle_momentum_z'].to_ndarray() / m_e / c
+    w  = ad['photons', 'particle_weight'].to_ndarray()
+
+    x_ind = ((x - ds.domain_left_edge[0].value) / dx[0]).astype(int)
+    y_ind = ((y - ds.domain_left_edge[1].value) / dx[1]).astype(int)
+    z_ind = ((z - ds.domain_left_edge[2].value) / dx[2]).astype(int)
+
+    zavg = np.zeros(ds.domain_dimensions)
+    uzavg = np.zeros(ds.domain_dimensions)
+    zuzavg = np.zeros(ds.domain_dimensions)
+    wavg = np.zeros(ds.domain_dimensions)
+
+    for i_p in range(len(x)):
+        zavg[x_ind[i_p],y_ind[i_p],z_ind[i_p]] += z[i_p] * w[i_p]
+        uzavg[x_ind[i_p],y_ind[i_p],z_ind[i_p]] += uz[i_p] * w[i_p]
+        zuzavg[x_ind[i_p],y_ind[i_p],z_ind[i_p]] += z[i_p] * uz[i_p] * w[i_p]
+        wavg[x_ind[i_p],y_ind[i_p],z_ind[i_p]] += w[i_p]
+
+    wavg_adj = np.where(wavg == 0, 1, wavg)
+    values_yt['photons: zavg'] = zavg / wavg_adj
+    values_yt['photons: uzavg'] = uzavg / wavg_adj
+    values_yt['photons: zuzavg'] = zuzavg / wavg_adj
+
+
+    values_rd = dict()
+    # Load reduced particle diagnostic data from plotfiles
+    values_rd['electrons: zavg'] = ad0[('boxlib','z_electrons')]
+    values_rd['protons: zavg'] = ad0[('boxlib','z_protons')]
+    values_rd['photons: zavg'] = ad0[('boxlib','z_photons')]
+
+    values_rd['electrons: uzavg'] = ad0[('boxlib','uz_electrons')]
+    values_rd['protons: uzavg'] = ad0[('boxlib','uz_protons')]
+    values_rd['photons: uzavg'] = ad0[('boxlib','uz_photons')]
+
+    values_rd['electrons: zuzavg'] = ad0[('boxlib','zuz_electrons')]
+    values_rd['protons: zuzavg'] = ad0[('boxlib','zuz_protons')]
+    values_rd['photons: zuzavg'] = ad0[('boxlib','zuz_photons')]
+
+    #--------------------------------------------------------------------------------------------------
+    # Part 3: compare values from plotfiles and diagnostics and print output
+    #--------------------------------------------------------------------------------------------------
+
+    error = dict()
+    tolerance = 5e-3 if single_precision else 1e-12
+
+    for k in values_yt.keys():
+        # check that the zeros line up, since we'll be ignoring them in the error calculation
+        assert(np.all((values_yt[k] == 0) == (values_rd[k] == 0)))
+        error[k] = np.max(abs(values_yt[k] - values_rd[k])[values_yt[k] != 0] / abs(values_yt[k])[values_yt[k] != 0])
+        print(k, 'relative error = ', error[k])
+        assert(error[k] < tolerance)
+
+    test_name = os.path.split(os.getcwd())[1]
+    checksumAPI.evaluate_checksum(test_name, fn)