1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
|
#!/usr/bin/env python3
#
# Copyright 2021 Remi Lehe
#
# This file is part of WarpX.
#
# License: BSD-3-Clause-LBNL
"""
This script tests the Gaussian-flux injection (and in particular
the rejection method in WarpX that we use to generate the right
velocity distribution).
Two population of particles are injected, with a slightly different
ratio of u_m/u_th. (This is in order to test the two different
rejection methods implemented in WarpX, which depend on the u_m/u_th ratio.)
After the particles are emitted with flux injection, this script produces
histograms of the velocity distribution and compares it with the expected
velocity distibution (Gaussian or Gaussian-flux depending on the direction
of space)
"""
import os
import re
import sys
import matplotlib.pyplot as plt
import numpy as np
from scipy.constants import c, m_e, m_p
from scipy.special import erf
import yt
sys.path.insert(1, '../../../../warpx/Regression/Checksum/')
import checksumAPI
yt.funcs.mylog.setLevel(0)
# Open plotfile specified in command line
fn = sys.argv[1]
ds = yt.load( fn )
ad = ds.all_data()
t_max = ds.current_time.item() # time of simulation
# Total number of electrons expected:
# Simulation parameters determine the total number of particles emitted (Ntot)
flux = 1. # in m^-2.s^-1, from the input script
emission_surface = 8*8 # in m^2
Ntot = flux * emission_surface * t_max
# Parameters of the histogram
hist_bins = 50
hist_range = [-0.5, 0.5]
# Define function that histogram and check the data
def gaussian_dist(u, u_th):
return 1./((2*np.pi)**.5*u_th) * np.exp(-u**2/(2*u_th**2) )
def gaussian_flux_dist(u, u_th, u_m):
au_m = np.abs(u_m)
normalization_factor = u_th**2 * np.exp(-au_m**2/(2*u_th**2)) + (np.pi/2)**.5*au_m*u_th * (1 + erf(au_m/(2**.5*u_th)))
result = 1./normalization_factor * np.where( u>0, u * np.exp(-(u-au_m)**2/(2*u_th**2)), 0 )
if u_m < 0.:
result = result[::-1]
return result
def compare_gaussian(u, w, u_th, label=''):
du = (hist_range[1]-hist_range[0])/hist_bins
w_hist, u_hist = np.histogram(u, bins=hist_bins, weights=w/du, range=hist_range)
u_hist = 0.5*(u_hist[1:]+u_hist[:-1])
w_th = Ntot*gaussian_dist(u_hist, u_th)
plt.plot( u_hist, w_hist, label=label+': simulation' )
plt.plot( u_hist, w_th, '--', label=label+': theory' )
assert np.allclose( w_hist, w_th, atol=0.07*w_th.max() )
def compare_gaussian_flux(u, w, u_th, u_m, label=''):
du = (hist_range[1]-hist_range[0])/hist_bins
w_hist, u_hist = np.histogram(u, bins=hist_bins, weights=w/du, range=hist_range)
u_hist = 0.5*(u_hist[1:]+u_hist[:-1])
w_th = Ntot*gaussian_flux_dist(u_hist, u_th, u_m)
plt.plot( u_hist, w_hist, label=label+': simulation' )
plt.plot( u_hist, w_th, '--', label=label+': theory' )
assert np.allclose( w_hist, w_th, atol=0.05*w_th.max() )
# Load data and perform check
plt.figure(figsize=(5,7))
plt.subplot(211)
plt.title('Electrons')
ux = ad['electron','particle_momentum_x'].to_ndarray()/(m_e*c)
uy = ad['electron','particle_momentum_y'].to_ndarray()/(m_e*c)
uz = ad['electron','particle_momentum_z'].to_ndarray()/(m_e*c)
w = ad['electron', 'particle_weight'].to_ndarray()
compare_gaussian(ux, w, u_th=0.1, label='u_x')
compare_gaussian_flux(uy, w, u_th=0.1, u_m=0.07, label='u_y')
compare_gaussian(uz, w, u_th=0.1, label='u_z')
plt.legend(loc=0)
plt.subplot(212)
plt.title('Protons')
ux = ad['proton','particle_momentum_x'].to_ndarray()/(m_p*c)
uy = ad['proton','particle_momentum_y'].to_ndarray()/(m_p*c)
uz = ad['proton','particle_momentum_z'].to_ndarray()/(m_p*c)
w = ad['proton', 'particle_weight'].to_ndarray()
compare_gaussian_flux(ux, w, u_th=0.1, u_m=-0.05, label='u_x')
compare_gaussian(uy, w, u_th=0.1, label='u_y')
compare_gaussian(uz, w, u_th=0.1, label='u_z')
plt.legend(loc=0)
plt.savefig('Distribution.png')
# Verify checksum
test_name = os.path.split(os.getcwd())[1]
if re.search( 'single_precision', fn ):
checksumAPI.evaluate_checksum(test_name, fn, rtol=1.e-3)
else:
checksumAPI.evaluate_checksum(test_name, fn)
|
