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
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
|
#!/usr/bin/env python3
#
# --- Test if MPI_communicator is correctly passed from Python to CPP
# --- This test should be run with MPI enabled
# --- Inputs taken from langmuir_2d. Runs 10 steps, and will check
# --- if the correct amount of processors are initialized in AMReX.
from mpi4py import MPI
from pywarpx import picmi
constants = picmi.constants
##########################
# MPI communicator setup
##########################
# split processor 0 into separate communicator from others
comm_world = MPI.COMM_WORLD
rank = comm_world.Get_rank()
if rank == 0:
color = 0
else:
color = 1
new_comm = comm_world.Split(color)
##########################
# physics parameters
##########################
# different communicators will be passed different plasma density
plasma_density = [1.e18, 1.e19]
plasma_xmin = 0.
plasma_x_velocity = 0.1*constants.c
##########################
# numerics parameters
##########################
# --- Number of time steps
max_steps = 10
diagnostic_intervals = "::10"
# --- Grid
nx = 64
ny = 64
xmin = -20.e-6
ymin = -20.e-6
xmax = +20.e-6
ymax = +20.e-6
number_per_cell_each_dim = [2,2]
##########################
# physics components
##########################
uniform_plasma = picmi.UniformDistribution(density = plasma_density[color],
upper_bound = [0., None, None],
directed_velocity = [0.1*constants.c, 0., 0.])
electrons = picmi.Species(particle_type='electron', name='electrons', initial_distribution=uniform_plasma)
##########################
# numerics components
##########################
grid = picmi.Cartesian2DGrid(number_of_cells = [nx, ny],
lower_bound = [xmin, ymin],
upper_bound = [xmax, ymax],
lower_boundary_conditions = ['periodic', 'periodic'],
upper_boundary_conditions = ['periodic', 'periodic'],
moving_window_velocity = [0., 0., 0.],
warpx_max_grid_size = 32)
solver = picmi.ElectromagneticSolver(grid=grid, cfl=1.)
##########################
# diagnostics
##########################
field_diag = picmi.FieldDiagnostic(name = f'diag{color + 1}',
grid = grid,
period = diagnostic_intervals,
data_list = ['Ex', 'Jx'],
write_dir = '.',
warpx_file_prefix = f'Python_pass_mpi_comm_plt{color + 1}_')
part_diag = picmi.ParticleDiagnostic(name = f'diag{color + 1}',
period = diagnostic_intervals,
species = [electrons],
data_list = ['weighting', 'ux'])
##########################
# simulation setup
##########################
sim = picmi.Simulation(solver = solver,
max_steps = max_steps,
verbose = 1,
warpx_current_deposition_algo = 'direct')
sim.add_species(electrons,
layout = picmi.GriddedLayout(n_macroparticle_per_cell=number_per_cell_each_dim, grid=grid))
sim.add_diagnostic(field_diag)
sim.add_diagnostic(part_diag)
##########################
# simulation run
##########################
sim.step(max_steps, mpi_comm=new_comm)
##########################
# test
##########################
# NOTE: Some tests are done in this input PICMI file. These tests
# check that amrex initialized the correct amount of procs and
# that the procs ranks in amrex are correct.
# If any of these tests fail, the terminal will print that the
# program crashed.
comm_world_size = comm_world.size
new_comm_size = new_comm.size
if color == 0:
# verify that communicator contains correct number of procs (1)
assert sim.extension.getNProcs() == comm_world_size - 1
assert sim.extension.getNProcs() == new_comm_size
else:
# verify that amrex initialized with 1 fewer proc than comm world
assert sim.extension.getNProcs() == comm_world_size - 1
assert sim.extension.getNProcs() == new_comm_size
# verify that amrex proc ranks are offset by -1 from
# world comm proc ranks
assert sim.extension.getMyProc() == rank - 1
|
