Correct and test fusion module in RZ geometry (#3255)

* Add particle rotation in NuclearFusionFunc.H

* Minor

* indent

* initial work

* fixed bugs and added species

* update documentation

* delete unused file

* Add properties for neutron, hydrogen isotopes, helium isotopes

* Update code to be more consistent

* Correct typo

* Parse deuterium-tritium fusion

* Start putting in place the files for deuterium-tritium

* Update documentation

* Prepare structures for deuterium tritium

* Fix typo

* Fix compilation

* Add neutron

* Add correct formula for the cross-section

* Correct compilation error

* Fix nuclear fusion

* Reset benchmarks

* Prepare creation functor for 2-product fusion

* First implementation of momentum initialization

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

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

* Use utility function for fusion

* Minor modification of variable names

* Fix GPU compilation

* Fix single precision compilation

* Update types

* Use util function in P-B fusion

* Correct compilation errors

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

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

* Correct errors

* Update values of mass and charge

* Correct compilation error

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

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

* Correct compilation error

* Correct compilation error

* Correct compilation error

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

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

* Reset benchmark

* Use helium particle in proton-boron, to avoid resetting benchmark

* Fixed proton-boron test

* Revert "Fixed proton-boron test"

This reverts commit 73c8d9d0be8417d5cd08a23daeebbc322c984808.

* Incorporate Neil's recommendations

* Reset benchmarks

* Correct compilation errors

* Add new deuterium tritium automated test

* Correct formula of cross-section

* Correct cross-section

* Improve analysis script

* Add test of energy conservation

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

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

* Add test of conservation of momentum

* Progress in analysis script

* Fix error in the initial energy of the deuterium particles

* Add check of isotropy

* Clean up the test script

* Rewrite p_sq formula in a way to avoids machine-precision negative numbers

* Add checksum

* Clean up code

* Add test for fusion in RZ geometry

* Update code to take into account actual timestep

* Update RZ test

* Update RZ test

* Increase number of particles

* Impart radial memory on DT particles

* Correct RZ momenta

* Remove unused file

* Update test

* Fix definition of theta

* Add new test

* Add checksum

* Update test

* Update tests

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

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

* Fix Python analysis script

* Remove CPU and ID from new benchmark

Co-authored-by: Yinjian Zhao <yinjianzhao@lbl.gov>
Co-authored-by: Luca Fedeli <luca.fedeli@cea.fr>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

2 files changed, 165 insertions, 15 deletions

diff --git a/Examples/Modules/nuclear_fusion/analysis_deuterium_tritium_fusion.py b/Examples/Modules/nuclear_fusion/analysis_deuterium_tritium_fusion.py
index f218df54d..b5f5da683 100755
--- a/Examples/Modules/nuclear_fusion/analysis_deuterium_tritium_fusion.py
+++ b/Examples/Modules/nuclear_fusion/analysis_deuterium_tritium_fusion.py
@@ -6,6 +6,7 @@
 # License: BSD-3-Clause-LBNL
 
 import os
+import re
 import sys
 
 import yt
@@ -66,15 +67,30 @@ barn_to_square_meter = 1.e-28
 
 E_fusion = 17.5893*MeV_to_Joule # Energy released during the fusion reaction
 
+## Checks whether this is the 2D or the 3D test
+warpx_used_inputs = open('./warpx_used_inputs', 'r').read()
+if re.search('geometry.dims = RZ', warpx_used_inputs):
+    is_RZ = True
+else:
+    is_RZ = False
+
 ## Some numerical parameters for this test
 size_x = 8
 size_y = 8
 size_z = 16
-dV_total = size_x*size_y*size_z # Total simulation volume
+if is_RZ:
+    dV_slice = np.pi * size_x**2
+    yt_z_string = "particle_position_y"
+    nppcell_1 = 10000*8
+    nppcell_2 = 900*8
+else:
+    dV_slice = size_x*size_y
+    yt_z_string = "particle_position_z"
+    nppcell_1 = 10000
+    nppcell_2 = 900
 # Volume of a slice corresponding to a single cell in the z direction. In tests 1 and 2, all the
 # particles of a given species in the same slice have the exact same momentum
-dV_slice = size_x*size_y
-dt = 1./(scc.c*np.sqrt(3.))
+
 # In test 1 and 2, the energy in cells number i (in z direction) is typically Energy_step * i**2
 Energy_step = 22.*keV_to_Joule
 
@@ -89,7 +105,7 @@ def add_existing_species_to_dict(yt_ad, data_dict, species_name, prefix, suffix)
     data_dict[prefix+"_w_"+suffix]   = yt_ad[species_name, "particle_weight"].v
     data_dict[prefix+"_id_"+suffix]  = yt_ad[species_name, "particle_id"].v
     data_dict[prefix+"_cpu_"+suffix] = yt_ad[species_name, "particle_cpu"].v
-    data_dict[prefix+"_z_"+suffix]   = yt_ad[species_name, "particle_position_z"].v
+    data_dict[prefix+"_z_"+suffix]   = yt_ad[species_name, yt_z_string].v
 
 def add_empty_species_to_dict(data_dict, species_name, prefix, suffix):
     data_dict[prefix+"_px_"+suffix]  = np.empty(0)
@@ -263,8 +279,11 @@ def expected_weight_com(E_com, reactant0_density, reactant1_density, dV, dt):
 def check_macroparticle_number(data, fusion_probability_target_value, num_pair_per_cell):
     ## Checks that the number of macroparticles is as expected for the first and second tests
 
-    ## The first slice 0 < z < 1 does not contribute to product species creation
-    numcells = dV_total - dV_slice
+    ## The first slice 0 < z < 1 does not contribute to alpha creation
+    if is_RZ:
+        numcells = size_x*(size_z-1)
+    else:
+        numcells = size_x*size_y*(size_z-1)
     ## In these tests, the fusion_multiplier is so high that the fusion probability per pair is
     ## equal to the parameter fusion_probability_target_value
     fusion_probability_per_pair = fusion_probability_target_value
@@ -315,7 +334,7 @@ def compute_E_com2(data):
     p_reactant0_sq = 2.*mass[reactant_species[0]]*(Energy_step*np.arange(size_z)**2)
     return p_sq_reactant1_frame_to_E_COM_frame(p_reactant0_sq)
 
-def check_fusion_yield(data, expected_fusion_number, E_com, reactant0_density, reactant1_density):
+def check_fusion_yield(data, expected_fusion_number, E_com, reactant0_density, reactant1_density, dt):
     ## Checks that the fusion yield is as expected for the first and second tests.
     product_weight_theory = expected_weight_com(E_com/keV_to_Joule,
         reactant0_density, reactant1_density, dV_slice, dt)
@@ -330,24 +349,25 @@ def check_fusion_yield(data, expected_fusion_number, E_com, reactant0_density, r
         assert(np.all(is_close(product_weight_theory, product_weight_simulation,
                                rtol = 5.*relative_std_weight)))
 
-def specific_check1(data):
-    check_isotropy(data, relative_tolerance = 3.e-2)
+def specific_check1(data, dt):
+    if not is_RZ:
+        check_isotropy(data, relative_tolerance = 3.e-2)
     expected_fusion_number = check_macroparticle_number(data,
                                                         fusion_probability_target_value = 0.002,
-                                                        num_pair_per_cell = 10000)
+                                                        num_pair_per_cell = nppcell_1)
     E_com = compute_E_com1(data)
     check_fusion_yield(data, expected_fusion_number, E_com, reactant0_density = 1.,
-                                                           reactant1_density = 1.)
+                                                           reactant1_density = 1., dt=dt)
 
-def specific_check2(data):
+def specific_check2(data, dt):
     check_xy_isotropy(data)
     ## Only 900 particles pairs per cell here because we ignore the 10% of reactants that are at rest
     expected_fusion_number = check_macroparticle_number(data,
                                                         fusion_probability_target_value = 0.02,
-                                                        num_pair_per_cell = 900)
+                                                        num_pair_per_cell = nppcell_2)
     E_com = compute_E_com2(data)
     check_fusion_yield(data, expected_fusion_number, E_com, reactant0_density = 1.e20,
-                                                           reactant1_density = 1.e26)
+                                                           reactant1_density = 1.e26, dt=dt)
 
 def check_charge_conservation(rho_start, rho_end):
     assert(np.all(is_close(rho_start, rho_end, rtol=2.e-11)))
@@ -359,6 +379,7 @@ def main():
     ds_start = yt.load(filename_start)
     ad_end = ds_end.all_data()
     ad_start = ds_start.all_data()
+    dt = float(ds_end.current_time - ds_start.current_time)
     field_data_end = ds_end.covering_grid(level=0, left_edge=ds_end.domain_left_edge,
                                           dims=ds_end.domain_dimensions)
     field_data_start = ds_start.covering_grid(level=0, left_edge=ds_start.domain_left_edge,
@@ -379,7 +400,7 @@ def main():
         generic_check(data)
 
         # Checks that are specific to test number i
-        eval("specific_check"+str(i)+"(data)")
+        eval("specific_check"+str(i)+"(data, dt)")
 
     rho_start = field_data_start["rho"].to_ndarray()
     rho_end = field_data_end["rho"].to_ndarray()
diff --git a/Examples/Modules/nuclear_fusion/inputs_deuterium_tritium_rz b/Examples/Modules/nuclear_fusion/inputs_deuterium_tritium_rz
new file mode 100644
index 000000000..fb581c825
--- /dev/null
+++ b/Examples/Modules/nuclear_fusion/inputs_deuterium_tritium_rz
@@ -0,0 +1,129 @@
+#################################
+####### GENERAL PARAMETERS ######
+#################################
+## With these parameters, each cell has a size of exactly 1 by 1 by 1
+max_step = 1
+amr.n_cell = 8 16
+amr.max_grid_size = 8
+amr.blocking_factor = 8
+amr.max_level = 0
+geometry.dims = RZ
+geometry.prob_lo     = 0.    0.
+geometry.prob_hi     = 8.   16.
+
+#################################
+###### Boundary Condition #######
+#################################
+boundary.field_lo = none periodic
+boundary.field_hi = pec periodic
+
+#################################
+############ NUMERICS ###########
+#################################
+warpx.verbose = 1
+warpx.cfl = 1.0
+
+# Order of particle shape factors
+algo.particle_shape = 1
+
+#################################
+############ PLASMA #############
+#################################
+particles.species_names = deuterium1 tritium1 helium1 neutron1 deuterium2 tritium2 helium2 neutron2
+
+my_constants.m_deuterium = 2.01410177812*m_u
+my_constants.m_tritium = 3.0160492779*m_u
+my_constants.m_reduced = m_deuterium*m_tritium/(m_deuterium+m_tritium)
+my_constants.keV_to_J = 1.e3*q_e
+my_constants.Energy_step = 22. * keV_to_J
+
+deuterium1.species_type = deuterium
+deuterium1.injection_style = "NRandomPerCell"
+deuterium1.num_particles_per_cell = 80000
+deuterium1.profile = constant
+deuterium1.density = 1.
+deuterium1.momentum_distribution_type = parse_momentum_function
+## Thanks to the floor, all particles in the same cell have the exact same momentum
+deuterium1.momentum_function_ux(x,y,z) = "u = sqrt(2*m_reduced*Energy_step*(floor(z)**2))/(m_deuterium*clight); if(x*x+y*y>0.0, -u*y/sqrt(x*x+y*y), 0.0)" # azimuthal velocity
+deuterium1.momentum_function_uy(x,y,z) = "u = sqrt(2*m_reduced*Energy_step*(floor(z)**2))/(m_deuterium*clight); if(x*x+y*y>0.0, u*x/sqrt(x*x+y*y), 0.0)" # azimuthal velocity
+deuterium1.momentum_function_uz(x,y,z) = "0"
+deuterium1.do_not_push = 1
+deuterium1.do_not_deposit = 1
+
+tritium1.species_type = tritium
+tritium1.injection_style = "NRandomPerCell"
+tritium1.num_particles_per_cell = 80000
+tritium1.profile = constant
+tritium1.density = 1.
+tritium1.momentum_distribution_type = "parse_momentum_function"
+## Thanks to the floor, all particles in the same cell have the exact same momentum
+tritium1.momentum_function_ux(x,y,z) = "u = sqrt(2*m_reduced*Energy_step*(floor(z)**2))/(m_tritium*clight); if(x*x+y*y>0.0, u*y/sqrt(x*x+y*y), 0.0)" # counter-streaming azimuthal velocity
+tritium1.momentum_function_uy(x,y,z) = "u = sqrt(2*m_reduced*Energy_step*(floor(z)**2))/(m_tritium*clight); if(x*x+y*y>0.0, -u*x/sqrt(x*x+y*y), 0.0)" # counter-streaming azimuthal velocity
+tritium1.momentum_function_uz(x,y,z) = 0
+tritium1.do_not_push = 1
+tritium1.do_not_deposit = 1
+
+helium1.species_type = helium4
+helium1.do_not_push = 1
+helium1.do_not_deposit = 1
+
+neutron1.species_type = neutron
+neutron1.do_not_push = 1
+neutron1.do_not_deposit = 1
+
+my_constants.background_dens = 1.e26
+my_constants.beam_dens = 1.e20
+
+deuterium2.species_type = deuterium
+deuterium2.injection_style = "NRandomPerCell"
+deuterium2.num_particles_per_cell = 8000
+deuterium2.profile = "parse_density_function"
+## A tenth of the macroparticles in each cell is made of immobile high-density background deuteriums.
+## The other nine tenths are made of fast low-density beam deuteriums.
+deuterium2.density_function(x,y,z) = if(y - floor(y) < 0.1, 10.*background_dens, 10./9.*beam_dens)
+deuterium2.momentum_distribution_type = "parse_momentum_function"
+deuterium2.momentum_function_ux(x,y,z) = 0.
+deuterium2.momentum_function_uy(x,y,z) = 0.
+deuterium2.momentum_function_uz(x,y,z) = "if(y - floor(y) < 0.1,
+                                          0., sqrt(2*m_deuterium*Energy_step*(floor(z)**2))/(m_deuterium*clight))"
+deuterium2.do_not_push = 1
+deuterium2.do_not_deposit = 1
+
+tritium2.species_type = tritium
+tritium2.injection_style = "NRandomPerCell"
+tritium2.num_particles_per_cell = 800
+tritium2.profile = constant
+tritium2.density = background_dens
+tritium2.momentum_distribution_type = "constant"
+tritium2.do_not_push = 1
+tritium2.do_not_deposit = 1
+
+helium2.species_type = helium4
+helium2.do_not_push = 1
+helium2.do_not_deposit = 1
+
+neutron2.species_type = neutron
+neutron2.do_not_push = 1
+neutron2.do_not_deposit = 1
+
+#################################
+############ COLLISION ##########
+#################################
+collisions.collision_names = DTF1 DTF2
+
+DTF1.species = deuterium1 tritium1
+DTF1.product_species = helium1 neutron1
+DTF1.type = nuclearfusion
+DTF1.fusion_multiplier = 1.e50
+
+DTF2.species = deuterium2 tritium2
+DTF2.product_species = helium2 neutron2
+DTF2.type = nuclearfusion
+DTF2.fusion_multiplier = 1.e15
+DTF2.fusion_probability_target_value = 0.02
+
+# Diagnostics
+diagnostics.diags_names = diag1
+diag1.intervals = 1
+diag1.diag_type = Full
+diag1.fields_to_plot = rho