diff options
27 files changed, 1797 insertions, 598 deletions
diff --git a/Examples/Tests/Langmuir/inputs.multi.rzmodes.rt b/Examples/Tests/Langmuir/inputs.multi.rzmodes.rt new file mode 100644 index 000000000..9808163f4 --- /dev/null +++ b/Examples/Tests/Langmuir/inputs.multi.rzmodes.rt @@ -0,0 +1,95 @@ +# Runs test with RZ multimode solver +# Maximum number of time steps +max_step = 40 + +# number of grid points +amr.n_cell = 64 200 + +# Maximum allowable size of each subdomain in the problem domain; +# this is used to decompose the domain for parallel calculations. +amr.max_grid_size = 32 + +# Maximum level in hierarchy (for now must be 0, i.e., one level in total) +amr.max_level = 0 + +amr.plot_int = 40 # How often to write plotfiles. "<= 0" means no plotfiles. + +# Geometry +geometry.coord_sys = 1 # 0: Cartesian +geometry.is_periodic = 0 1 # Is periodic? +geometry.prob_lo = 0.e-6 -20.e-6 # physical domain +geometry.prob_hi = 20.e-6 20.e-6 + +warpx.nmodes = 3 + +warpx.serialize_ics = 1 +warpx.plot_raw_fields = 1 + +# Verbosity +warpx.verbose = 1 + +# Interpolation +interpolation.nox = 1 +interpolation.noy = 1 +interpolation.noz = 1 + +# CFL +warpx.cfl = 1.0 + +# dive is not supported with RZ multimode +warpx.do_dive_cleaning = 0 + +# Parameters for the plasma wave +my_constants.epsilon0 = 0.001 +my_constants.epsilon1 = 0.001 +my_constants.epsilon2 = 0.001 +my_constants.kp = 266125.0928256017 +my_constants.k0 = 314159.2653589793 +my_constants.w0 = 5.e-6 +# Note: kp is calculated in SI for a density of 2e24 +# k0 is calculated so as to have 2 periods within the 40e-6 wide box. + +# Particles +particles.nspecies = 2 +particles.species_names = electrons ions + +electrons.charge = -q_e +electrons.mass = m_e +electrons.injection_style = "NUniformPerCell" +electrons.num_particles_per_cell_each_dim = 2 8 2 +electrons.xmin = 0.e-6 +electrons.xmax = 18.e-6 +electrons.zmin = -20.e-6 +electrons.zmax = +20.e-6 + +electrons.profile = constant +electrons.density = 2.e24 # number of electrons per m^3 +electrons.momentum_distribution_type = parse_momentum_function +electrons.momentum_function_ux(x,y,z) = "+ epsilon0/kp*2*x/w0**2*exp(-(x**2+y**2)/w0**2)*sin(k0*z) + - epsilon1/kp*2/w0*exp(-(x**2+y**2)/w0**2)*sin(k0*z) + + epsilon1/kp*4*x**2/w0**3*exp(-(x**2+y**2)/w0**2)*sin(k0*z) + - epsilon2/kp*8*x/w0**2*exp(-(x**2+y**2)/w0**2)*sin(k0*z) + + epsilon2/kp*8*x*(x**2-y**2)/w0**4*exp(-(x**2+y**2)/w0**2)*sin(k0*z)" +electrons.momentum_function_uy(x,y,z) = "+ epsilon0/kp*2*y/w0**2*exp(-(x**2+y**2)/w0**2)*sin(k0*z) + + epsilon1/kp*4*x*y/w0**3*exp(-(x**2+y**2)/w0**2)*sin(k0*z) + + epsilon2/kp*8*y/w0**2*exp(-(x**2+y**2)/w0**2)*sin(k0*z) + + epsilon2/kp*8*y*(x**2-y**2)/w0**4*exp(-(x**2+y**2)/w0**2)*sin(k0*z)" +electrons.momentum_function_uz(x,y,z) = "- epsilon0/kp*k0*exp(-(x**2+y**2)/w0**2)*cos(k0*z) + - epsilon1/kp*k0*2*x/w0*exp(-(x**2+y**2)/w0**2)*cos(k0*z) + - epsilon2/kp*k0*4*(x**2-y**2)/w0**2*exp(-(x**2+y**2)/w0**2)*cos(k0*z)" + +ions.charge = q_e +ions.mass = m_p +ions.injection_style = "NUniformPerCell" +ions.num_particles_per_cell_each_dim = 2 8 2 +ions.xmin = 0.e-6 +ions.xmax = 18.e-6 +ions.zmin = -20.e-6 +ions.zmax = +20.e-6 + +ions.profile = constant +ions.density = 2.e24 # number of ions per m^3 +ions.momentum_distribution_type = constant +ions.ux = 0. +ions.uy = 0. +ions.uz = 0. diff --git a/Examples/Tests/Langmuir/langmuir_PICMI_rz_multimode_analyze.py b/Examples/Tests/Langmuir/langmuir_PICMI_rz_multimode_analyze.py new file mode 100644 index 000000000..9d9fa40f6 --- /dev/null +++ b/Examples/Tests/Langmuir/langmuir_PICMI_rz_multimode_analyze.py @@ -0,0 +1,170 @@ +# This is a script that analyses the multimode simulation results. +# This simulates a RZ multimode periodic plasma wave. +# The electric field from the simulation is compared to the analytic value + +import numpy as np +import matplotlib +matplotlib.use('Agg') +import matplotlib.pyplot as plt +from pywarpx import picmi + +nr = 64 +nz = 200 + +rmin = 0.e0 +zmin = 0.e0 +rmax = +20.e-6 +zmax = +40.e-6 + +# Parameters describing particle distribution +density = 2.e24 +epsilon0 = 0.001 +epsilon1 = 0.001 +epsilon2 = 0.001 +w0 = 5.e-6 +n_osc_z = 3 + +# Wave vector of the wave +k0 = 2.*np.pi*n_osc_z/(zmax - zmin) + +# Plasma frequency +wp = np.sqrt((density*picmi.q_e**2)/(picmi.m_e*picmi.ep0)) +kp = wp/picmi.c + +uniform_plasma = picmi.UniformDistribution(density = density, + upper_bound = [+18e-6, +18e-6, None], + directed_velocity = [0., 0., 0.]) + +momentum_expressions = ["""+ epsilon0/kp*2*x/w0**2*exp(-(x**2+y**2)/w0**2)*sin(k0*z) + - epsilon1/kp*2/w0*exp(-(x**2+y**2)/w0**2)*sin(k0*z) + + epsilon1/kp*4*x**2/w0**3*exp(-(x**2+y**2)/w0**2)*sin(k0*z) + - epsilon2/kp*8*x/w0**2*exp(-(x**2+y**2)/w0**2)*sin(k0*z) + + epsilon2/kp*8*x*(x**2-y**2)/w0**4*exp(-(x**2+y**2)/w0**2)*sin(k0*z)""", + """+ epsilon0/kp*2*y/w0**2*exp(-(x**2+y**2)/w0**2)*sin(k0*z) + + epsilon1/kp*4*x*y/w0**3*exp(-(x**2+y**2)/w0**2)*sin(k0*z) + + epsilon2/kp*8*y/w0**2*exp(-(x**2+y**2)/w0**2)*sin(k0*z) + + epsilon2/kp*8*y*(x**2-y**2)/w0**4*exp(-(x**2+y**2)/w0**2)*sin(k0*z)""", + """- epsilon0/kp*k0*exp(-(x**2+y**2)/w0**2)*cos(k0*z) + - epsilon1/kp*k0*2*x/w0*exp(-(x**2+y**2)/w0**2)*cos(k0*z) + - epsilon2/kp*k0*4*(x**2-y**2)/w0**2*exp(-(x**2+y**2)/w0**2)*cos(k0*z)"""] + +analytic_plasma = picmi.AnalyticDistribution(density_expression = density, + upper_bound = [+18e-6, +18e-6, None], + epsilon0 = epsilon0, + epsilon1 = epsilon1, + epsilon2 = epsilon2, + kp = kp, + k0 = k0, + w0 = w0, + momentum_expressions = momentum_expressions) + +electrons = picmi.Species(particle_type='electron', name='electrons', initial_distribution=analytic_plasma) +ions = picmi.Species(particle_type='proton', name='protons', initial_distribution=uniform_plasma) + +grid = picmi.CylindricalGrid(number_of_cells = [nr, nz], + n_azimuthal_modes = 3, + lower_bound = [rmin, zmin], + upper_bound = [rmax, zmax], + lower_boundary_conditions = ['dirichlet', 'periodic'], + upper_boundary_conditions = ['dirichlet', 'periodic'], + moving_window_velocity = [0., 0.], + warpx_max_grid_size=64) + +solver = picmi.ElectromagneticSolver(grid=grid, cfl=1.) + +sim = picmi.Simulation(solver = solver, + max_steps = 40, + verbose = 1, + warpx_plot_int = 40, + warpx_current_deposition_algo = 'direct', + warpx_field_gathering_algo = 'standard', + warpx_particle_pusher_algo = 'boris') + +sim.add_species(electrons, layout=picmi.GriddedLayout(n_macroparticle_per_cell=[2,8,2], grid=grid)) +sim.add_species(ions, layout=picmi.GriddedLayout(n_macroparticle_per_cell=[2,8,2], grid=grid)) + +# write_inputs will create an inputs file that can be used to run +# with the compiled version. +#sim.write_input_file(file_name='inputsrz_from_PICMI') + +# Alternatively, sim.step will run WarpX, controlling it from Python +sim.step() + + +# Below is WarpX specific code to check the results. +import pywarpx +from pywarpx.fields import * + +def calcEr( z, r, k0, w0, wp, t, epsilons) : + """ + Return the radial electric field as an array + of the same length as z and r, in the half-plane theta=0 + """ + Er_array = ( + epsilons[0] * picmi.m_e*picmi.c**2/picmi.q_e * 2*r/w0**2 * + np.exp( -r**2/w0**2 ) * np.sin( k0*z ) * np.sin( wp*t ) + - epsilons[1] * picmi.m_e*picmi.c**2/picmi.q_e * 2/w0 * + np.exp( -r**2/w0**2 ) * np.sin( k0*z ) * np.sin( wp*t ) + + epsilons[1] * picmi.m_e*picmi.c**2/picmi.q_e * 4*r**2/w0**3 * + np.exp( -r**2/w0**2 ) * np.sin( k0*z ) * np.sin( wp*t ) + - epsilons[2] * picmi.m_e*picmi.c**2/picmi.q_e * 8*r/w0**2 * + np.exp( -r**2/w0**2 ) * np.sin( k0*z ) * np.sin( wp*t ) + + epsilons[2] * picmi.m_e*picmi.c**2/picmi.q_e * 8*r**3/w0**4 * + np.exp( -r**2/w0**2 ) * np.sin( k0*z ) * np.sin( wp*t )) + return( Er_array ) + +def calcEz( z, r, k0, w0, wp, t, epsilons) : + """ + Return the longitudinal electric field as an array + of the same length as z and r, in the half-plane theta=0 + """ + Ez_array = ( + - epsilons[0] * picmi.m_e*picmi.c**2/picmi.q_e * k0 * + np.exp( -r**2/w0**2 ) * np.cos( k0*z ) * np.sin( wp*t ) + - epsilons[1] * picmi.m_e*picmi.c**2/picmi.q_e * k0 * 2*r/w0 * + np.exp( -r**2/w0**2 ) * np.cos( k0*z ) * np.sin( wp*t ) + - epsilons[2] * picmi.m_e*picmi.c**2/picmi.q_e * k0 * 4*r**2/w0**2 * + np.exp( -r**2/w0**2 ) * np.cos( k0*z ) * np.sin( wp*t )) + return( Ez_array ) + +# Get node centered coordinates +dr = (rmax - rmin)/nr +dz = (zmax - zmin)/nz +coords = np.indices([nr+1, nz+1], 'd') +rr = rmin + coords[0]*dr +zz = zmin + coords[1]*dz + +# Current time of the simulation +t0 = pywarpx._libwarpx.libwarpx.warpx_gett_new(0) + +# Get the raw field data. Note that these are the real and imaginary +# parts of the fields for each azimuthal mode. +Ex_sim_modes = ExWrapper()[...] +Ez_sim_modes = EzWrapper()[...] + +# Sum the real components to get the field along x-axis (theta = 0) +Er_sim = np.sum(Ex_sim_modes[:,:,::2], axis=2) +Ez_sim = np.sum(Ez_sim_modes[:,:,::2], axis=2) + +# The analytical solutions +Er_th = calcEr(zz[:-1,:], rr[:-1,:] + dr/2., k0, w0, wp, t0, [epsilon0, epsilon1, epsilon2]) +Ez_th = calcEz(zz[:,:-1] + dz/2., rr[:,:-1], k0, w0, wp, t0, [epsilon0, epsilon1, epsilon2]) + +max_error_Er = abs(Er_sim - Er_th).max()/abs(Er_th).max() +max_error_Ez = abs(Ez_sim - Ez_th).max()/abs(Ez_th).max() +print("Max error Er %e"%max_error_Er) +print("Max error Ez %e"%max_error_Er) + +# Plot the last field from the loop (Ez at iteration 40) +plt.subplot2grid( (1,2), (0,0) ) +plt.imshow( Ez_sim ) +plt.colorbar() +plt.title('Ez, last iteration\n(simulation)') +plt.subplot2grid( (1,2), (0,1) ) +plt.imshow( Ez_th ) +plt.colorbar() +plt.title('Ez, last iteration\n(theory)') +plt.tight_layout() +plt.savefig('langmuir_multi_rz_multimode_analysis.png') + +assert max(max_error_Er, max_error_Ez) < 0.02 diff --git a/Examples/Tests/Langmuir/langmuir_multi_rz_multimode_analysis.py b/Examples/Tests/Langmuir/langmuir_multi_rz_multimode_analysis.py new file mode 100644 index 000000000..b87dd0fae --- /dev/null +++ b/Examples/Tests/Langmuir/langmuir_multi_rz_multimode_analysis.py @@ -0,0 +1,124 @@ +#! /usr/bin/env python + +# This is a script that analyses the simulation results from +# the script `inputs.multi.rzmodes.rt`. This simulates a RZ multimode periodic plasma wave. +# The electric field from the simulation is compared to the analytic value +import sys +import matplotlib +matplotlib.use('Agg') +import matplotlib.pyplot as plt +import yt +yt.funcs.mylog.setLevel(50) +import numpy as np +from scipy.constants import e, m_e, epsilon_0, c + +# this will be the name of the plot file +fn = sys.argv[1] + +# Parameters (these parameters must match the parameters in `inputs.multi.rzmm.rt`) +epsilon0 = 0.001 +epsilon1 = 0.001 +epsilon2 = 0.001 +n = 2.e24 +w0 = 5.e-6 +n_osc_z = 2 +gscale = 1 +rmin = 0e-6; rmax = 20.e-6; Nr = 64//gscale +zmin = -20e-6; zmax = 20.e-6; Nz = 200//gscale + +# Wave vector of the wave +k0 = 2.*np.pi*n_osc_z/(zmax-zmin) +# Plasma frequency +wp = np.sqrt((n*e**2)/(m_e*epsilon_0)) +kp = wp/c + +def Jx( z, r, k0, w0, wp, t) : + x = r + y = 0 + return -n*e*c*(+ epsilon0 /kp * 2*x/w0**2 * np.exp( -r**2/w0**2 ) * np.sin( k0*z ) * np.cos( wp*t ) + - epsilon1 /kp * 2/w0 * np.exp( -r**2/w0**2 ) * np.sin( k0*z ) * np.cos( wp*t ) + + epsilon1 /kp * 4*x**2/w0**3 * np.exp( -r**2/w0**2 ) * np.sin( k0*z ) * np.cos( wp*t ) + - epsilon2 /kp * 8*x/w0**2 * np.exp( -r**2/w0**2 ) * np.sin( k0*z ) * np.cos( wp*t ) + + epsilon2 /kp * 8*x*(x**2-y**2)/w0**4 * np.exp( -r**2/w0**2 ) * np.sin( k0*z ) * np.cos( wp*t )) + +def Jz( z, r, k0, w0, wp, t) : + x = r + y = 0 + return +n*e*c*(- epsilon0 /kp * k0 * np.exp( -r**2/w0**2 ) * np.cos( k0*z ) * np.cos( wp*t ) + - epsilon1 /kp * k0 * 2*x/w0 * np.exp( -r**2/w0**2 ) * np.cos( k0*z ) * np.cos( wp*t ) + - epsilon2 /kp * k0 * 4*(x**2-y**2)/w0**2 * np.exp( -r**2/w0**2 ) * np.cos( k0*z ) * np.cos( wp*t )) + +def Er( z, r, k0, w0, wp, t) : + """ + Return the radial electric field as an array + of the same length as z and r, in the half-plane theta=0 + """ + Er_array = ( + epsilon0 * m_e*c**2/e * 2*r/w0**2 * + np.exp( -r**2/w0**2 ) * np.sin( k0*z ) * np.sin( wp*t ) + - epsilon1 * m_e*c**2/e * 2/w0 * + np.exp( -r**2/w0**2 ) * np.sin( k0*z ) * np.sin( wp*t ) + + epsilon1 * m_e*c**2/e * 4*r**2/w0**3 * + np.exp( -r**2/w0**2 ) * np.sin( k0*z ) * np.sin( wp*t ) + - epsilon2 * m_e*c**2/e * 8*r/w0**2 * + np.exp( -r**2/w0**2 ) * np.sin( k0*z ) * np.sin( wp*t ) + + epsilon2 * m_e*c**2/e * 8*r**3/w0**4 * + np.exp( -r**2/w0**2 ) * np.sin( k0*z ) * np.sin( wp*t )) + return( Er_array ) + +def Ez( z, r, k0, w0, wp, t) : + """ + Return the longitudinal electric field as an array + of the same length as z and r, in the half-plane theta=0 + """ + Ez_array = ( + - epsilon0 * m_e*c**2/e * k0 * + np.exp( -r**2/w0**2 ) * np.cos( k0*z ) * np.sin( wp*t ) + - epsilon1 * m_e*c**2/e * k0 * 2*r/w0 * + np.exp( -r**2/w0**2 ) * np.cos( k0*z ) * np.sin( wp*t ) + - epsilon2 * m_e*c**2/e * k0 * 4*r**2/w0**2 * + np.exp( -r**2/w0**2 ) * np.cos( k0*z ) * np.sin( wp*t )) + return( Ez_array ) + +# Read the file +ds = yt.load(fn) +t0 = ds.current_time.to_ndarray().mean() +data = ds.covering_grid(level=0, left_edge=ds.domain_left_edge, + dims=ds.domain_dimensions) + +# Get cell centered coordinates +dr = (rmax - rmin)/Nr +dz = (zmax - zmin)/Nz +coords = np.indices([Nr, Nz], 'd') +rr = rmin + (coords[0] + 0.5)*dr +zz = zmin + (coords[1] + 0.5)*dz + +# Check the validity of the fields +overall_max_error = 0 +Er_sim = data['Ex'].to_ndarray()[:,:,0] +Er_th = Er(zz, rr, k0, w0, wp, t0) +max_error = abs(Er_sim-Er_th).max()/abs(Er_th).max() +print('Er: Max error: %.2e' %(max_error)) +overall_max_error = max( overall_max_error, max_error ) + +Ez_sim = data['Ez'].to_ndarray()[:,:,0] +Ez_th = Ez(zz, rr, k0, w0, wp, t0) +max_error = abs(Ez_sim-Ez_th).max()/abs(Ez_th).max() +print('Ez: Max error: %.2e' %(max_error)) +overall_max_error = max( overall_max_error, max_error ) + +# Plot the last field from the loop (Ez at iteration 40) +plt.subplot2grid( (1,2), (0,0) ) +plt.imshow( Ez_sim ) +plt.colorbar() +plt.title('Ez, last iteration\n(simulation)') +plt.subplot2grid( (1,2), (0,1) ) +plt.imshow( Ez_th ) +plt.colorbar() +plt.title('Ez, last iteration\n(theory)') +plt.tight_layout() +plt.savefig('langmuir_multi_rz_analysis.png') + +# Automatically check the validity +assert overall_max_error < 0.02 + diff --git a/Python/pywarpx/PGroup.py b/Python/pywarpx/PGroup.py index 68b37740d..48e68ceb5 100644 --- a/Python/pywarpx/PGroup.py +++ b/Python/pywarpx/PGroup.py @@ -83,6 +83,10 @@ class PGroup(object): return _libwarpx.get_particle_y(self.ispecie)[self.igroup] yp = property(getyp) + def getrp(self): + return _libwarpx.get_particle_r(self.ispecie)[self.igroup] + rp = property(getrp) + def getzp(self): return _libwarpx.get_particle_z(self.ispecie)[self.igroup] zp = property(getzp) @@ -136,6 +140,10 @@ class PGroup(object): return _libwarpx.get_particle_Bz(self.ispecie)[self.igroup] bz = property(getbz) + def gettheta(self): + return _libwarpx.get_particle_theta(self.ispecie)[self.igroup] + theta = property(gettheta) + class PGroups(object): def __init__(self, ispecie=0): self.ispecie = ispecie diff --git a/Python/pywarpx/WarpInterface.py b/Python/pywarpx/WarpInterface.py new file mode 100644 index 000000000..f84c22a5d --- /dev/null +++ b/Python/pywarpx/WarpInterface.py @@ -0,0 +1,171 @@ +import warp +from . import fields +from pywarpx import PGroup + + +def warp_species(warp_type, picmie_species): + """Returns a Warp species that has a reference to the WarpX particles. + """ + elec_pgroups = PGroup.PGroups(ispecie=picmie_species.species_number) + return warp.Species(type=warp_type, pgroups=elec_pgroups) + + +class _WarpX_FIELDtype(object): + """Mirrors part of the EM3D_FIELDtype type from Warp + yf + """ + def __init__(self, yf): + self.yf = yf + + +class _WarpX_BLOCKtype(object): + """Mirrors part of the EM3D_BLOCKtype type from Warp + core + xmin, ymin, zmin + xmax, ymax, zmax + dx, dy, dz + xrbnd, yrbnd, zrbnd + """ + def __init__(self, fields, picmi_grid): + self.core = _WarpX_FIELDtype(fields) + self.picmi_grid = picmi_grid + + self.xmin = self.picmi_grid.lower_bound[0] + if self.picmi_grid.number_of_dimensions == 3: + self.ymin = self.picmi_grid.lower_bound[1] + else: + self.ymin = 0. + self.zmin = self.picmi_grid.lower_bound[-1] + + self.xmax = self.picmi_grid.upper_bound[0] + if self.picmi_grid.number_of_dimensions == 3: + self.ymax = self.picmi_grid.upper_bound[1] + else: + self.ymax = 0. + self.zmax = self.picmi_grid.upper_bound[-1] + + self.dx = (self.xmax - self.xmin)/self.picmi_grid.number_of_cells[0] + if self.picmi_grid.number_of_dimensions == 3: + self.dy = (self.ymax - self.ymin)/self.picmi_grid.number_of_cells[1] + else: + self.dy = 1. + self.dz = (self.zmax - self.zmin)/self.picmi_grid.number_of_cells[-1] + + self.xrbnd = 0 + self.yrbnd = 0 + self.zrbnd = 0 + + +class _WarpX_YEEFIELDtype(object): + """Mirrors part of the EM3D_YEEFIELDtype type from Warp + Exp, Eyp, Ezp + Bxp, Byp, Bzp + Ex, Ey, Ez + Bx, By, Bz + Jx, Jy, Jz + """ + def __init__(self, level=0): + self.level = level + self._Ex_wrap = fields.ExWrapper(level, include_ghosts=True) + self._Ey_wrap = fields.EyWrapper(level, include_ghosts=True) + self._Ez_wrap = fields.EzWrapper(level, include_ghosts=True) + self._Bx_wrap = fields.BxWrapper(level, include_ghosts=True) + self._By_wrap = fields.ByWrapper(level, include_ghosts=True) + self._Bz_wrap = fields.BzWrapper(level, include_ghosts=True) + self._Jx_wrap = fields.JxWrapper(level, include_ghosts=True) + self._Jy_wrap = fields.JyWrapper(level, include_ghosts=True) + self._Jz_wrap = fields.JzWrapper(level, include_ghosts=True) + + self._Ex_wrap._getlovects() # --- Calculated nghosts + self.nxguard = self._Ex_wrap.nghosts + self.nyguard = self._Ex_wrap.nghosts + self.nzguard = self._Ex_wrap.nghosts + + def _get_wrapped_array(self, wrapper): + result = wrapper[...] + if len(result.shape) == 2: + # --- Add the middle dimension that Warp is expecting + result.shape = [result.shape[0], 1, result.shape[1]] + return result + + @property + def Exp(self): + return self._get_wrapped_array(self._Ex_wrap) + @property + def Eyp(self): + return self._get_wrapped_array(self._Ey_wrap) + @property + def Ezp(self): + return self._get_wrapped_array(self._Ez_wrap) + @property + def Bxp(self): + return self._get_wrapped_array(self._Bx_wrap) + @property + def Byp(self): + return self._get_wrapped_array(self._By_wrap) + @property + def Bzp(self): + return self._get_wrapped_array(self._Bz_wrap) + @property + def Ex(self): + return self._get_wrapped_array(self._Ex_wrap) + @property + def Ey(self): + return self._get_wrapped_array(self._Ey_wrap) + @property + def Ez(self): + return self._get_wrapped_array(self._Ez_wrap) + @property + def Bx(self): + return self._get_wrapped_array(self._Bx_wrap) + @property + def By(self): + return self._get_wrapped_array(self._By_wrap) + @property + def Bz(self): + return self._get_wrapped_array(self._Bz_wrap) + @property + def Jx(self): + return self._get_wrapped_array(self._Jx_wrap) + @property + def Jy(self): + return self._get_wrapped_array(self._Jy_wrap) + @property + def Jz(self): + return self._get_wrapped_array(self._Jz_wrap) + + +class WarpX_EM3D(warp.EM3D): + """Mirrors part of the Warp EM3D class, mostly diagnostics. + """ + def __init__(self, picmi_grid, level=0): + self.picmi_grid = picmi_grid + self.level = level + + # --- Only define what is necessary for the diagnostics + self.fields = _WarpX_YEEFIELDtype(level) + self.block = _WarpX_BLOCKtype(self.fields, picmi_grid) + + self.isactive = True + + self.l_1dz = (picmi_grid.number_of_dimensions == 1) + self.l_2dxz = (picmi_grid.number_of_dimensions == 2) + try: + picmi_grid.nr + except AttributeError: + self.l_2drz = False + else: + self.l_2drz = True + + self.l4symtry = False + self.l2symtry = False + + self.nx = picmi_grid.number_of_cells[0] + if not self.l_2dxz: + self.ny = picmi_grid.number_of_cells[1] + else: + self.ny = 0 + self.nz = picmi_grid.number_of_cells[-1] + + self.zgrid = 0. # --- This should be obtained from WarpX + diff --git a/Python/pywarpx/_libwarpx.py b/Python/pywarpx/_libwarpx.py index 4c3283b97..2a37dbd5e 100755 --- a/Python/pywarpx/_libwarpx.py +++ b/Python/pywarpx/_libwarpx.py @@ -1,6 +1,7 @@ # --- This defines the wrapper functions that directly call the underlying compiled routines import os import sys +import atexit import ctypes from ctypes.util import find_library as _find_library import numpy as np @@ -8,6 +9,13 @@ from numpy.ctypeslib import ndpointer as _ndpointer from .Geometry import geometry +try: + # --- If mpi4py is going to be used, this needs to be imported + # --- before libwarpx is loaded (though don't know why) + from mpi4py import MPI +except ImportError: + pass + # --- Is there a better way of handling constants? clight = 2.99792458e+8 # m/s @@ -184,6 +192,7 @@ def initialize(argv=None): libwarpx.warpx_init() +@atexit.register def finalize(finalize_mpi=1): ''' @@ -399,7 +408,10 @@ def get_particle_x(species_number): ''' structs = get_particle_structs(species_number) - return [struct['x'] for struct in structs] + if geometry_dim == '3d' or geometry_dim == '2d': + return [struct['x'] for struct in structs] + elif geometry_dim == 'rz': + return [struct['x']*np.cos(theta) for struct, theta in zip(structs, get_particle_theta(species_number))] def get_particle_y(species_number): @@ -410,7 +422,26 @@ def get_particle_y(species_number): ''' structs = get_particle_structs(species_number) - return [struct['y'] for struct in structs] + if geometry_dim == '3d' or geometry_dim == '2d': + return [struct['y'] for struct in structs] + elif geometry_dim == 'rz': + return [struct['x']*np.sin(theta) for struct, theta in zip(structs, get_particle_theta(species_number))] + + +def get_particle_r(species_number): + ''' + + Return a list of numpy arrays containing the particle 'r' + positions on each tile. + + ''' + structs = get_particle_structs(species_number) + if geometry_dim == 'rz': + return [struct['x'] for struct in structs] + elif geometry_dim == '3d': + return [np.sqrt(struct['x']**2 + struct['y']**2) for struct in structs] + elif geometry_dim == '2d': + raise Exception('get_particle_r: There is no r coordinate with 2D Cartesian') def get_particle_z(species_number): @@ -556,6 +587,53 @@ def get_particle_Bz(species_number): return get_particle_arrays(species_number, 9) +def get_particle_theta(species_number): + ''' + + Return a list of numpy arrays containing the particle + theta on each tile. + + ''' + + if geometry_dim == 'rz': + return get_particle_arrays(species_number, 10) + elif geometry_dim == '3d': + return [np.arctan2(struct['y'], struct['x']) for struct in structs] + elif geometry_dim == '2d': + raise Exception('get_particle_r: There is no theta coordinate with 2D Cartesian') + + +def _get_mesh_field_list(warpx_func, level, direction, include_ghosts): + """ + Generic routine to fetch the list of field data arrays. + """ + shapes = _LP_c_int() + size = ctypes.c_int(0) + ncomps = ctypes.c_int(0) + ngrow = ctypes.c_int(0) + data = warpx_func(level, direction, + ctypes.byref(size), ctypes.byref(ncomps), + ctypes.byref(ngrow), ctypes.byref(shapes)) + ng = ngrow.value + grid_data = [] + shapesize = dim + if ncomps.value > 1: + shapesize += 1 + for i in range(size.value): + shape = tuple([shapes[shapesize*i + d] for d in range(shapesize)]) + # --- The data is stored in Fortran order, hence shape is reversed and a transpose is taken. + arr = np.ctypeslib.as_array(data[i], shape[::-1]).T + arr.setflags(write=1) + if include_ghosts: + grid_data.append(arr) + else: + grid_data.append(arr[tuple([slice(ng, -ng) for _ in range(dim)])]) + + _libc.free(shapes) + _libc.free(data) + return grid_data + + def get_mesh_electric_field(level, direction, include_ghosts=True): ''' @@ -582,28 +660,7 @@ def get_mesh_electric_field(level, direction, include_ghosts=True): ''' assert(level == 0) - - shapes = _LP_c_int() - size = ctypes.c_int(0) - ngrow = ctypes.c_int(0) - data = libwarpx.warpx_getEfield(level, direction, - ctypes.byref(size), ctypes.byref(ngrow), - ctypes.byref(shapes)) - ng = ngrow.value - grid_data = [] - for i in range(size.value): - shape = tuple([shapes[dim*i + d] for d in range(dim)]) - # --- The data is stored in Fortran order, hence shape is reversed and a transpose is taken. - arr = np.ctypeslib.as_array(data[i], shape[::-1]).T - arr.setflags(write=1) - if include_ghosts: - grid_data.append(arr) - else: - grid_data.append(arr[[slice(ng, -ng) for _ in range(dim)]]) - - _libc.free(shapes) - _libc.free(data) - return grid_data + return _get_mesh_field_list(libwarpx.warpx_getEfield, level, direction, include_ghosts) def get_mesh_electric_field_cp(level, direction, include_ghosts=True): @@ -631,28 +688,7 @@ def get_mesh_electric_field_cp(level, direction, include_ghosts=True): ''' assert(level == 0) - - shapes = _LP_c_int() - size = ctypes.c_int(0) - ngrow = ctypes.c_int(0) - data = libwarpx.warpx_getEfieldCP(level, direction, - ctypes.byref(size), ctypes.byref(ngrow), - ctypes.byref(shapes)) - ng = ngrow.value - grid_data = [] - for i in range(size.value): - shape = tuple([shapes[dim*i + d] for d in range(dim)]) - # --- The data is stored in Fortran order, hence shape is reversed and a transpose is taken. - arr = np.ctypeslib.as_array(data[i], shape[::-1]).T - arr.setflags(write=1) - if include_ghosts: - grid_data.append(arr) - else: - grid_data.append(arr[[slice(ng, -ng) for _ in range(dim)]]) - - _libc.free(shapes) - _libc.free(data) - return grid_data + return _get_mesh_field_list(libwarpx.warpx_getEfieldCP, level, direction, include_ghosts) def get_mesh_electric_field_fp(level, direction, include_ghosts=True): @@ -680,28 +716,7 @@ def get_mesh_electric_field_fp(level, direction, include_ghosts=True): ''' assert(level == 0) - - shapes = _LP_c_int() - size = ctypes.c_int(0) - ngrow = ctypes.c_int(0) - data = libwarpx.warpx_getEfieldFP(level, direction, - ctypes.byref(size), ctypes.byref(ngrow), - ctypes.byref(shapes)) - ng = ngrow.value - grid_data = [] - for i in range(size.value): - shape = tuple([shapes[dim*i + d] for d in range(dim)]) - # --- The data is stored in Fortran order, hence shape is reversed and a transpose is taken. - arr = np.ctypeslib.as_array(data[i], shape[::-1]).T - arr.setflags(write=1) - if include_ghosts: - grid_data.append(arr) - else: - grid_data.append(arr[[slice(ng, -ng) for _ in range(dim)]]) - - _libc.free(shapes) - _libc.free(data) - return grid_data + return _get_mesh_field_list(libwarpx.warpx_getEfieldFP, level, direction, include_ghosts) def get_mesh_magnetic_field(level, direction, include_ghosts=True): @@ -730,28 +745,7 @@ def get_mesh_magnetic_field(level, direction, include_ghosts=True): ''' assert(level == 0) - - shapes = _LP_c_int() - size = ctypes.c_int(0) - ngrow = ctypes.c_int(0) - data = libwarpx.warpx_getBfield(level, direction, - ctypes.byref(size), ctypes.byref(ngrow), - ctypes.byref(shapes)) - ng = ngrow.value - grid_data = [] - for i in range(size.value): - shape = tuple([shapes[dim*i + d] for d in range(dim)]) - # --- The data is stored in Fortran order, hence shape is reversed and a transpose is taken. - arr = np.ctypeslib.as_array(data[i], shape[::-1]).T - arr.setflags(write=1) - if include_ghosts: - grid_data.append(arr) - else: - grid_data.append(arr[[slice(ng, -ng) for _ in range(dim)]]) - - _libc.free(shapes) - _libc.free(data) - return grid_data + return _get_mesh_field_list(libwarpx.warpx_getBfield, level, direction, include_ghosts) def get_mesh_magnetic_field_cp(level, direction, include_ghosts=True): @@ -779,28 +773,7 @@ def get_mesh_magnetic_field_cp(level, direction, include_ghosts=True): ''' assert(level == 0) - - shapes = _LP_c_int() - size = ctypes.c_int(0) - ngrow = ctypes.c_int(0) - data = libwarpx.warpx_getBfieldCP(level, direction, - ctypes.byref(size), ctypes.byref(ngrow), - ctypes.byref(shapes)) - ng = ngrow.value - grid_data = [] - for i in range(size.value): - shape = tuple([shapes[dim*i + d] for d in range(dim)]) - # --- The data is stored in Fortran order, hence shape is reversed and a transpose is taken. - arr = np.ctypeslib.as_array(data[i], shape[::-1]).T - arr.setflags(write=1) - if include_ghosts: - grid_data.append(arr) - else: - grid_data.append(arr[[slice(ng, -ng) for _ in range(dim)]]) - - _libc.free(shapes) - _libc.free(data) - return grid_data + return _get_mesh_field_list(libwarpx.warpx_getBfieldCP, level, direction, include_ghosts) def get_mesh_magnetic_field_fp(level, direction, include_ghosts=True): @@ -828,28 +801,7 @@ def get_mesh_magnetic_field_fp(level, direction, include_ghosts=True): ''' assert(level == 0) - - shapes = _LP_c_int() - size = ctypes.c_int(0) - ngrow = ctypes.c_int(0) - data = libwarpx.warpx_getBfieldFP(level, direction, - ctypes.byref(size), ctypes.byref(ngrow), - ctypes.byref(shapes)) - ng = ngrow.value - grid_data = [] - for i in range(size.value): - shape = tuple([shapes[dim*i + d] for d in range(dim)]) - # --- The data is stored in Fortran order, hence shape is reversed and a transpose is taken. - arr = np.ctypeslib.as_array(data[i], shape[::-1]).T - arr.setflags(write=1) - if include_ghosts: - grid_data.append(arr) - else: - grid_data.append(arr[[slice(ng, -ng) for _ in range(dim)]]) - - _libc.free(shapes) - _libc.free(data) - return grid_data + return _get_mesh_field_list(libwarpx.warpx_getBfieldFP, level, direction, include_ghosts) def get_mesh_current_density(level, direction, include_ghosts=True): @@ -876,28 +828,7 @@ def get_mesh_current_density(level, direction, include_ghosts=True): ''' assert(level == 0) - - shapes = _LP_c_int() - size = ctypes.c_int(0) - ngrow = ctypes.c_int(0) - data = libwarpx.warpx_getCurrentDensity(level, direction, - ctypes.byref(size), ctypes.byref(ngrow), - ctypes.byref(shapes)) - ng = ngrow.value - grid_data = [] - for i in range(size.value): - shape = tuple([shapes[dim*i + d] for d in range(dim)]) - # --- The data is stored in Fortran order, hence shape is reversed and a transpose is taken. - arr = np.ctypeslib.as_array(data[i], shape[::-1]).T - arr.setflags(write=1) - if include_ghosts: - grid_data.append(arr) - else: - grid_data.append(arr[[slice(ng, -ng) for _ in range(dim)]]) - - _libc.free(shapes) - _libc.free(data) - return grid_data + return _get_mesh_field_list(libwarpx.warpx_getCurrentDensity, level, direction, include_ghosts) def get_mesh_current_density_cp(level, direction, include_ghosts=True): @@ -925,28 +856,7 @@ def get_mesh_current_density_cp(level, direction, include_ghosts=True): ''' assert(level == 0) - - shapes = _LP_c_int() - size = ctypes.c_int(0) - ngrow = ctypes.c_int(0) - data = libwarpx.warpx_getCurrentDensityCP(level, direction, - ctypes.byref(size), ctypes.byref(ngrow), - ctypes.byref(shapes)) - ng = ngrow.value - grid_data = [] - for i in range(size.value): - shape = tuple([shapes[dim*i + d] for d in range(dim)]) - # --- The data is stored in Fortran order, hence shape is reversed and a transpose is taken. - arr = np.ctypeslib.as_array(data[i], shape[::-1]).T - arr.setflags(write=1) - if include_ghosts: - grid_data.append(arr) - else: - grid_data.append(arr[[slice(ng, -ng) for _ in range(dim)]]) - - _libc.free(shapes) - _libc.free(data) - return grid_data + return _get_mesh_field_list(libwarpx.warpx_getCurrentDensityCP, level, direction, include_ghosts) def get_mesh_current_density_fp(level, direction, include_ghosts=True): @@ -974,28 +884,7 @@ def get_mesh_current_density_fp(level, direction, include_ghosts=True): ''' assert(level == 0) - - shapes = _LP_c_int() - size = ctypes.c_int(0) - ngrow = ctypes.c_int(0) - data = libwarpx.warpx_getCurrentDensityFP(level, direction, - ctypes.byref(size), ctypes.byref(ngrow), - ctypes.byref(shapes)) - ng = ngrow.value - grid_data = [] - for i in range(size.value): - shape = tuple([shapes[dim*i + d] for d in range(dim)]) - # --- The data is stored in Fortran order, hence shape is reversed and a transpose is taken. - arr = np.ctypeslib.as_array(data[i], shape[::-1]).T - arr.setflags(write=1) - if include_ghosts: - grid_data.append(arr) - else: - grid_data.append(arr[[slice(ng, -ng) for _ in range(dim)]]) - - _libc.free(shapes) - _libc.free(data) - return grid_data + return _get_mesh_field_list(libwarpx.warpx_getCurrentDensityFP, level, direction, include_ghosts) def _get_mesh_array_lovects(level, direction, include_ghosts=True, getarrayfunc=None): diff --git a/Python/pywarpx/fields.py b/Python/pywarpx/fields.py index 8237ed867..9068a22ea 100644 --- a/Python/pywarpx/fields.py +++ b/Python/pywarpx/fields.py @@ -11,7 +11,7 @@ import numpy as np from . import _libwarpx -class MultiFABWrapper(object): +class _MultiFABWrapper(object): """Wrapper around field arrays at level 0 This provides a convenient way to query and set fields that are broken up into FABs. The indexing is based on global indices. @@ -21,16 +21,23 @@ class MultiFABWrapper(object): - get_fabs: routine that returns the list of FABs - level=0: ignored """ - def __init__(self, direction, overlaps, get_lovects, get_fabs, level=0): + def __init__(self, direction, overlaps, get_lovects, get_fabs, level=0, include_ghosts=False): self.direction = direction self.overlaps = np.array(overlaps) self.get_lovects = get_lovects self.get_fabs = get_fabs self.level = 0 #level - self.include_ghosts = False + self.include_ghosts = include_ghosts + + self.dim = _libwarpx.dim + if self.dim == 2: + # --- Grab the first and last overlaps (for x and z) + self.overlaps = self.overlaps[::2] def _getlovects(self): - return self.get_lovects(self.level, self.direction, self.include_ghosts) + lovects = self.get_lovects(self.level, self.direction, self.include_ghosts) + self.nghosts = -lovects.min() + return lovects def _gethivects(self): lovects = self._getlovects() @@ -38,7 +45,7 @@ class MultiFABWrapper(object): hivects = np.zeros_like(lovects) for i in range(len(fields)): - hivects[:,i] = lovects[:,i] + np.array(fields[i].shape) - self.overlaps + hivects[:,i] = lovects[:,i] + np.array(fields[i].shape[:self.dim]) - self.overlaps return hivects @@ -50,14 +57,29 @@ class MultiFABWrapper(object): def __getitem__(self, index): """Returns slices of a decomposed array, The shape of - the object returned depends on the number of ix, iy and iz specified, which - can be from none to all three. Note that the values of ix, iy and iz are - relative to the fortran indexing, meaning that 0 is the lower boundary - of the domain. + the object returned depends on the number of ix, iy and iz specified, which + can be from none to all three. Note that the values of ix, iy and iz are + relative to the fortran indexing, meaning that 0 is the lower boundary + of the whole domain. """ if index == Ellipsis: - index = (slice(None), slice(None), slice(None)) - + index = tuple(self.dim*[slice(None)]) + + if len(index) < self.dim: + # --- Add extra dims to index if needed + index = list(index) + for i in range(len(index), self.dim): + index.append(slice(None)) + index = tuple(index) + + if self.dim == 2: + return self._getitem2d(index) + elif self.dim == 3: + return self._getitem3d(index) + + def _getitem3d(self, index): + """Returns slices of a 3D decomposed array, + """ ix = index[0] iy = index[1] iz = index[2] @@ -66,25 +88,25 @@ class MultiFABWrapper(object): hivects = self._gethivects() fields = self._getfields() - nx = hivects[0,:].max() - ny = hivects[1,:].max() - nz = hivects[2,:].max() + nx = hivects[0,:].max() - self.nghosts + ny = hivects[1,:].max() - self.nghosts + nz = hivects[2,:].max() - self.nghosts if isinstance(ix, slice): - ixstart = max(ix.start, 0) - ixstop = min(ix.stop or nx + 1, nx + self.overlaps[0]) + ixstart = max(ix.start or -self.nghosts, -self.nghosts) + ixstop = min(ix.stop or nx + 1 + self.nghosts, nx + self.overlaps[0] + self.nghosts) else: ixstart = ix ixstop = ix + 1 if isinstance(iy, slice): - iystart = max(iy.start, 0) - iystop = min(iy.stop or ny + 1, ny + self.overlaps[1]) + iystart = max(iy.start or -self.nghosts, -self.nghosts) + iystop = min(iy.stop or ny + 1 + self.nghosts, ny + self.overlaps[1] + self.nghosts) else: iystart = iy iystop = iy + 1 if isinstance(iz, slice): - izstart = max(iz.start, 0) - izstop = min(iz.stop or nz + 1, nz + self.overlaps[2]) + izstart = max(iz.start or -self.nghosts, -self.nghosts) + izstop = min(iz.stop or nz + 1 + self.nghosts, nz + self.overlaps[2] + self.nghosts) else: izstart = iz izstop = iz + 1 @@ -99,11 +121,11 @@ class MultiFABWrapper(object): # --- The ix1, 2 etc are relative to global indexing ix1 = max(ixstart, lovects[0,i]) - ix2 = min((ixstop or nx+1), lovects[0,i] + fields[i].shape[0]) + ix2 = min(ixstop, lovects[0,i] + fields[i].shape[0]) iy1 = max(iystart, lovects[1,i]) - iy2 = min((iystop or ny+1), lovects[1,i] + fields[i].shape[1]) + iy2 = min(iystop, lovects[1,i] + fields[i].shape[1]) iz1 = max(izstart, lovects[2,i]) - iz2 = min((izstop or nz+1), lovects[2,i] + fields[i].shape[2]) + iz2 = min(izstop, lovects[2,i] + fields[i].shape[2]) if ix1 < ix2 and iy1 < iy2 and iz1 < iz2: @@ -126,7 +148,84 @@ class MultiFABWrapper(object): if not isinstance(iz, slice): sss[2] = 0 - return resultglobal[sss] + return resultglobal[tuple(sss)] + + def _getitem2d(self, index): + """Returns slices of a 2D decomposed array, + """ + + lovects = self._getlovects() + hivects = self._gethivects() + fields = self._getfields() + + ix = index[0] + iz = index[1] + + if len(fields[0].shape) > self.dim: + ncomps = fields[0].shape[-1] + else: + ncomps = 1 + + if len(index) > self.dim: + if ncomps > 1: + ic = index[2] + else: + raise Exception('Too many indices given') + else: + ic = None + + nx = hivects[0,:].max() - self.nghosts + nz = hivects[1,:].max() - self.nghosts + + if isinstance(ix, slice): + ixstart = max(ix.start or -self.nghosts, -self.nghosts) + ixstop = min(ix.stop or nx + 1 + self.nghosts, nx + self.overlaps[0] + self.nghosts) + else: + ixstart = ix + ixstop = ix + 1 + if isinstance(iz, slice): + izstart = max(iz.start or -self.nghosts, -self.nghosts) + izstop = min(iz.stop or nz + 1 + self.nghosts, nz + self.overlaps[1] + self.nghosts) + else: + izstart = iz + izstop = iz + 1 + + # --- Setup the size of the array to be returned and create it. + # --- Space is added for multiple components if needed. + sss = (max(0, ixstop - ixstart), + max(0, izstop - izstart)) + if ncomps > 1 and ic is None: + sss = tuple(list(sss) + [ncomps]) + resultglobal = np.zeros(sss) + + for i in range(len(fields)): + + # --- The ix1, 2 etc are relative to global indexing + ix1 = max(ixstart, lovects[0,i]) + ix2 = min(ixstop, lovects[0,i] + fields[i].shape[0]) + iz1 = max(izstart, lovects[1,i]) + iz2 = min(izstop, lovects[1,i] + fields[i].shape[1]) + + if ix1 < ix2 and iz1 < iz2: + + sss = (slice(ix1 - lovects[0,i], ix2 - lovects[0,i]), + slice(iz1 - lovects[1,i], iz2 - lovects[1,i])) + if ic is not None: + sss = tuple(list(sss) + [ic]) + + vslice = (slice(ix1 - ixstart, ix2 - ixstart), + slice(iz1 - izstart, iz2 - izstart)) + + resultglobal[vslice] = fields[i][sss] + + # --- Now remove any of the reduced dimensions. + sss = [slice(None), slice(None)] + if not isinstance(ix, slice): + sss[0] = 0 + if not isinstance(iz, slice): + sss[1] = 0 + + return resultglobal[tuple(sss)] def __setitem__(self, index, value): """Sets slices of a decomposed array. The shape of @@ -135,8 +234,23 @@ class MultiFABWrapper(object): - value: input array (must be supplied) """ if index == Ellipsis: - index = (slice(None), slice(None), slice(None)) - + index = tuple(self.dim*[slice(None)]) + + if len(index) < self.dim: + # --- Add extra dims to index if needed + index = list(index) + for i in range(len(index), self.dim): + index.append(slice(None)) + index = tuple(index) + + if self.dim == 2: + return self._setitem2d(index, value) + elif self.dim == 3: + return self._setitem3d(index, value) + + def _setitem3d(self, index, value): + """Sets slices of a decomposed 3D array. + """ ix = index[0] iy = index[1] iz = index[2] @@ -145,9 +259,9 @@ class MultiFABWrapper(object): hivects = self._gethivects() fields = self._getfields() - nx = hivects[0,:].max() - ny = hivects[1,:].max() - nz = hivects[2,:].max() + nx = hivects[0,:].max() - self.nghosts + ny = hivects[1,:].max() - self.nghosts + nz = hivects[2,:].max() - self.nghosts # --- Add extra dimensions so that the input has the same number of # --- dimensions as array. @@ -160,20 +274,20 @@ class MultiFABWrapper(object): value3d.shape = sss if isinstance(ix, slice): - ixstart = max(ix.start, 0) - ixstop = min(ix.stop or nx + 1, nx + self.overlaps[0]) + ixstart = max(ix.start or -self.nghosts, -self.nghosts) + ixstop = min(ix.stop or nx + 1 + self.nghosts, nx + self.overlaps[0] + self.nghosts) else: ixstart = ix ixstop = ix + 1 if isinstance(iy, slice): - iystart = max(iy.start, 0) - iystop = min(iy.stop or ny + 1, ny + self.overlaps[1]) + iystart = max(iy.start or -self.nghosts, -self.nghosts) + iystop = min(iy.stop or ny + 1 + self.nghosts, ny + self.overlaps[1] + self.nghosts) else: iystart = iy iystop = iy + 1 if isinstance(iz, slice): - izstart = max(iz.start, 0) - izstop = min(iz.stop or nz + 1, nz + self.overlaps[2]) + izstart = max(iz.start or -self.nghosts, -self.nghosts) + izstop = min(iz.stop or nz + 1 + self.nghosts, nz + self.overlaps[2] + self.nghosts) else: izstart = iz izstop = iz + 1 @@ -182,11 +296,11 @@ class MultiFABWrapper(object): # --- The ix1, 2 etc are relative to global indexing ix1 = max(ixstart, lovects[0,i]) - ix2 = min((ixstop or nx+1), lovects[0,i] + fields[i].shape[0]) + ix2 = min(ixstop, lovects[0,i] + fields[i].shape[0]) iy1 = max(iystart, lovects[1,i]) - iy2 = min((iystop or ny+1), lovects[1,i] + fields[i].shape[1]) + iy2 = min(iystop, lovects[1,i] + fields[i].shape[1]) iz1 = max(izstart, lovects[2,i]) - iz2 = min((izstop or nz+1), lovects[2,i] + fields[i].shape[2]) + iz2 = min(izstop, lovects[2,i] + fields[i].shape[2]) if ix1 < ix2 and iy1 < iy2 and iz1 < iz2: @@ -202,49 +316,127 @@ class MultiFABWrapper(object): else: fields[i][sss] = value + def _setitem2d(self, index, value): + """Sets slices of a decomposed 2D array. + """ + ix = index[0] + iz = index[2] -def ExWrapper(level=0): - return MultiFABWrapper(direction=0, overlaps=[0,1,1], - get_lovects=_libwarpx.get_mesh_electric_field_lovects, - get_fabs=_libwarpx.get_mesh_electric_field, level=level) + lovects = self._getlovects() + hivects = self._gethivects() + fields = self._getfields() -def EyWrapper(level=0): - return MultiFABWrapper(direction=1, overlaps=[1,0,1], - get_lovects=_libwarpx.get_mesh_electric_field_lovects, - get_fabs=_libwarpx.get_mesh_electric_field, level=level) + if len(fields[0].shape) > self.dim: + ncomps = fields[0].shape[-1] + else: + ncomps = 1 -def EzWrapper(level=0): - return MultiFABWrapper(direction=2, overlaps=[1,1,0], - get_lovects=_libwarpx.get_mesh_electric_field_lovects, - get_fabs=_libwarpx.get_mesh_electric_field, level=level) + if len(index) > self.dim: + if ncomps > 1: + ic = index[2] + else: + raise Exception('Too many indices given') + else: + ic = None -def BxWrapper(level=0): - return MultiFABWrapper(direction=0, overlaps=[1,0,0], - get_lovects=_libwarpx.get_mesh_magnetic_field_lovects, - get_fabs=_libwarpx.get_mesh_magnetic_field, level=level) + nx = hivects[0,:].max() - self.nghosts + nz = hivects[2,:].max() - self.nghosts -def ByWrapper(level=0): - return MultiFABWrapper(direction=1, overlaps=[0,1,0], - get_lovects=_libwarpx.get_mesh_magnetic_field_lovects, - get_fabs=_libwarpx.get_mesh_magnetic_field, level=level) + # --- Add extra dimensions so that the input has the same number of + # --- dimensions as array. + if isinstance(value, np.ndarray): + value3d = np.array(value, copy=False) + sss = list(value3d.shape) + if not isinstance(ix, slice): sss[0:0] = [1] + if not isinstance(iz, slice): sss[1:1] = [1] + value3d.shape = sss -def BzWrapper(level=0): - return MultiFABWrapper(direction=2, overlaps=[0,0,1], - get_lovects=_libwarpx.get_mesh_magnetic_field_lovects, - get_fabs=_libwarpx.get_mesh_magnetic_field, level=level) + if isinstance(ix, slice): + ixstart = max(ix.start or -self.nghosts, -self.nghosts) + ixstop = min(ix.stop or nx + 1 + self.nghosts, nx + self.overlaps[0] + self.nghosts) + else: + ixstart = ix + ixstop = ix + 1 + if isinstance(iz, slice): + izstart = max(iz.start or -self.nghosts, -self.nghosts) + izstop = min(iz.stop or nz + 1 + self.nghosts, nz + self.overlaps[2] + self.nghosts) + else: + izstart = iz + izstop = iz + 1 + + for i in range(len(fields)): + + # --- The ix1, 2 etc are relative to global indexing + ix1 = max(ixstart, lovects[0,i]) + ix2 = min(ixstop, lovects[0,i] + fields[i].shape[0]) + iz1 = max(izstart, lovects[2,i]) + iz2 = min(izstop, lovects[2,i] + fields[i].shape[2]) + + if ix1 < ix2 and iz1 < iz2: -def JxWrapper(level=0): - return MultiFABWrapper(direction=0, overlaps=[0,1,1], - get_lovects=_libwarpx.get_mesh_current_density_lovects, - get_fabs=_libwarpx.get_mesh_current_density, level=level) + sss = (slice(ix1 - lovects[0,i], ix2 - lovects[0,i]), + slice(iz1 - lovects[2,i], iz2 - lovects[2,i])) + if ic is not None: + sss = tuple(list(sss) + [ic]) -def JyWrapper(level=0): - return MultiFABWrapper(direction=1, overlaps=[1,0,1], - get_lovects=_libwarpx.get_mesh_current_density_lovects, - get_fabs=_libwarpx.get_mesh_current_density, level=level) + if isinstance(value, np.ndarray): + vslice = (slice(ix1 - ixstart, ix2 - ixstart), + slice(iz1 - izstart, iz2 - izstart)) + fields[i][sss] = value3d[vslice] + else: + fields[i][sss] = value -def JzWrapper(level=0): - return MultiFABWrapper(direction=2, overlaps=[1,1,0], - get_lovects=_libwarpx.get_mesh_current_density_lovects, - get_fabs=_libwarpx.get_mesh_current_density, level=level) +def ExWrapper(level=0, include_ghosts=False): + return _MultiFABWrapper(direction=0, overlaps=[0,1,1], + get_lovects=_libwarpx.get_mesh_electric_field_lovects, + get_fabs=_libwarpx.get_mesh_electric_field, + level=level, include_ghosts=include_ghosts) + +def EyWrapper(level=0, include_ghosts=False): + return _MultiFABWrapper(direction=1, overlaps=[1,0,1], + get_lovects=_libwarpx.get_mesh_electric_field_lovects, + get_fabs=_libwarpx.get_mesh_electric_field, + level=level, include_ghosts=include_ghosts) + +def EzWrapper(level=0, include_ghosts=False): + return _MultiFABWrapper(direction=2, overlaps=[1,1,0], + get_lovects=_libwarpx.get_mesh_electric_field_lovects, + get_fabs=_libwarpx.get_mesh_electric_field, + level=level, include_ghosts=include_ghosts) + +def BxWrapper(level=0, include_ghosts=False): + return _MultiFABWrapper(direction=0, overlaps=[1,0,0], + get_lovects=_libwarpx.get_mesh_magnetic_field_lovects, + get_fabs=_libwarpx.get_mesh_magnetic_field, + level=level, include_ghosts=include_ghosts) + +def ByWrapper(level=0, include_ghosts=False): + return _MultiFABWrapper(direction=1, overlaps=[0,1,0], + get_lovects=_libwarpx.get_mesh_magnetic_field_lovects, + get_fabs=_libwarpx.get_mesh_magnetic_field, + level=level, include_ghosts=include_ghosts) + +def BzWrapper(level=0, include_ghosts=False): + return _MultiFABWrapper(direction=2, overlaps=[0,0,1], + get_lovects=_libwarpx.get_mesh_magnetic_field_lovects, + get_fabs=_libwarpx.get_mesh_magnetic_field, + level=level, include_ghosts=include_ghosts) + +def JxWrapper(level=0, include_ghosts=False): + return _MultiFABWrapper(direction=0, overlaps=[0,1,1], + get_lovects=_libwarpx.get_mesh_current_density_lovects, + get_fabs=_libwarpx.get_mesh_current_density, + level=level, include_ghosts=include_ghosts) + +def JyWrapper(level=0, include_ghosts=False): + return _MultiFABWrapper(direction=1, overlaps=[1,0,1], + get_lovects=_libwarpx.get_mesh_current_density_lovects, + get_fabs=_libwarpx.get_mesh_current_density, + level=level, include_ghosts=include_ghosts) + +def JzWrapper(level=0, include_ghosts=False): + return _MultiFABWrapper(direction=2, overlaps=[1,1,0], + get_lovects=_libwarpx.get_mesh_current_density_lovects, + get_fabs=_libwarpx.get_mesh_current_density, + level=level, include_ghosts=include_ghosts) diff --git a/Python/pywarpx/picmi.py b/Python/pywarpx/picmi.py index 3cbefc920..b89c3fc43 100644 --- a/Python/pywarpx/picmi.py +++ b/Python/pywarpx/picmi.py @@ -180,7 +180,6 @@ class AnalyticDistribution(picmistandard.PICMI_AnalyticDistribution): species.zmin = self.lower_bound[2] species.zmax = self.upper_bound[2] - # --- Only constant density is supported at this time species.profile = "parse_density_function" species.__setattr__('density_function(x,y,z)', self.density_expression) @@ -188,7 +187,10 @@ class AnalyticDistribution(picmistandard.PICMI_AnalyticDistribution): setattr(pywarpx.my_constants, k, v) # --- Note that WarpX takes gamma*beta as input - if np.any(np.not_equal(self.rms_velocity, 0.)): + if np.any(np.not_equal(self.momentum_expressions, None)): + species.momentum_distribution_type = 'parse_momentum_function' + self.setup_parse_momentum_functions(species) + elif np.any(np.not_equal(self.rms_velocity, 0.)): species.momentum_distribution_type = "gaussian" species.ux_m = self.directed_velocity[0]/c species.uy_m = self.directed_velocity[1]/c @@ -205,6 +207,19 @@ class AnalyticDistribution(picmistandard.PICMI_AnalyticDistribution): if self.fill_in: species.do_continuous_injection = 1 + def setup_parse_momentum_functions(self, species): + if self.momentum_expressions[0] is not None: + species.__setattr__('momentum_function_ux(x,y,z)', '({0})/{1}'.format(self.momentum_expressions[0], c)) + else: + species.__setattr__('momentum_function_ux(x,y,z)', '({0})/{1}'.format(self.directed_velocity[0], c)) + if self.momentum_expressions[1] is not None: + species.__setattr__('momentum_function_uy(x,y,z)', '({0})/{1}'.format(self.momentum_expressions[1], c)) + else: + species.__setattr__('momentum_function_uy(x,y,z)', '({0})/{1}'.format(self.directed_velocity[1], c)) + if self.momentum_expressions[2] is not None: + species.__setattr__('momentum_function_uz(x,y,z)', '({0})/{1}'.format(self.momentum_expressions[2], c)) + else: + species.__setattr__('momentum_function_uz(x,y,z)', '({0})/{1}'.format(self.directed_velocity[2], c)) class ParticleListDistribution(picmistandard.PICMI_ParticleListDistribution): def init(self, kw): @@ -259,13 +274,13 @@ class CylindricalGrid(picmistandard.PICMI_CylindricalGrid): assert self.lower_bound[0] >= 0., Exception('Lower radial boundary must be >= 0.') assert self.bc_rmin != 'periodic' and self.bc_rmax != 'periodic', Exception('Radial boundaries can not be periodic') - assert self.n_azimuthal_modes is None or self.n_azimuthal_modes == 1, Exception('Only one azimuthal mode supported') # Geometry pywarpx.geometry.coord_sys = 1 # RZ pywarpx.geometry.is_periodic = '0 %d'%(self.bc_zmin=='periodic') # Is periodic? pywarpx.geometry.prob_lo = self.lower_bound # physical domain pywarpx.geometry.prob_hi = self.upper_bound + pywarpx.warpx.nmodes = self.n_azimuthal_modes if self.moving_window_velocity is not None and np.any(np.not_equal(self.moving_window_velocity, 0.)): pywarpx.warpx.do_moving_window = 1 diff --git a/Source/Diagnostics/FieldIO.H b/Source/Diagnostics/FieldIO.H index 1a3b45580..24fd6abb6 100644 --- a/Source/Diagnostics/FieldIO.H +++ b/Source/Diagnostics/FieldIO.H @@ -15,6 +15,7 @@ PackPlotDataPtrs (Vector<const MultiFab*>& pmf, void AverageAndPackVectorField( MultiFab& mf_avg, const std::array< std::unique_ptr<MultiFab>, 3 >& vector_field, + const DistributionMapping& dm, const int dcomp, const int ngrow ); void diff --git a/Source/Diagnostics/FieldIO.cpp b/Source/Diagnostics/FieldIO.cpp index b1181f22f..0926a327e 100644 --- a/Source/Diagnostics/FieldIO.cpp +++ b/Source/Diagnostics/FieldIO.cpp @@ -189,6 +189,20 @@ WriteOpenPMDFields( const std::string& filename, } #endif // WARPX_USE_OPENPMD +void +ConstructTotalRZField(std::array< std::unique_ptr<MultiFab>, 3 >& mf_total, + const std::array< std::unique_ptr<MultiFab>, 3 >& vector_field) +{ + // Sum over the real components, giving quantity at theta=0 + MultiFab::Copy(*mf_total[0], *vector_field[0], 0, 0, 1, vector_field[0]->nGrowVect()); + MultiFab::Copy(*mf_total[1], *vector_field[1], 0, 0, 1, vector_field[1]->nGrowVect()); + MultiFab::Copy(*mf_total[2], *vector_field[2], 0, 0, 1, vector_field[2]->nGrowVect()); + for (int ic=2 ; ic < vector_field[0]->nComp() ; ic += 2) { + MultiFab::Add(*mf_total[0], *vector_field[0], ic, 0, 1, vector_field[0]->nGrowVect()); + MultiFab::Add(*mf_total[1], *vector_field[1], ic, 0, 1, vector_field[1]->nGrowVect()); + MultiFab::Add(*mf_total[2], *vector_field[2], ic, 0, 1, vector_field[2]->nGrowVect()); + } +} void PackPlotDataPtrs (Vector<const MultiFab*>& pmf, @@ -213,6 +227,7 @@ PackPlotDataPtrs (Vector<const MultiFab*>& pmf, void AverageAndPackVectorField( MultiFab& mf_avg, const std::array< std::unique_ptr<MultiFab>, 3 >& vector_field, + const DistributionMapping& dm, const int dcomp, const int ngrow ) { // The object below is temporary, and is needed because @@ -241,23 +256,89 @@ AverageAndPackVectorField( MultiFab& mf_avg, // - Face centered, in the same way as B on a Yee grid } else if ( vector_field[0]->is_nodal(0) ){ - PackPlotDataPtrs(srcmf, vector_field); - amrex::average_face_to_cellcenter( mf_avg, dcomp, srcmf, ngrow); + // Note that average_face_to_cellcenter operates only on the number of + // arrays equal to the number of dimensions. So, for 2D, PackPlotDataPtrs + // packs in the x and z (or r and z) arrays, which are then cell averaged. + // The Copy code then copies the z from the 2nd to the 3rd field, + // and copies over directly the y (or theta) component (which is + // already cell centered). + if (vector_field[0]->nComp() > 1) { + // When there are more than one components, the total + // fields needs to be constructed in temporary MultiFabs + // Note that mf_total is declared in the same way as + // vector_field so that it can be passed into PackPlotDataPtrs. + std::array<std::unique_ptr<MultiFab>,3> mf_total; + mf_total[0].reset(new MultiFab(vector_field[0]->boxArray(), dm, 1, vector_field[0]->nGrowVect())); + mf_total[1].reset(new MultiFab(vector_field[1]->boxArray(), dm, 1, vector_field[1]->nGrowVect())); + mf_total[2].reset(new MultiFab(vector_field[2]->boxArray(), dm, 1, vector_field[2]->nGrowVect())); + ConstructTotalRZField(mf_total, vector_field); + PackPlotDataPtrs(srcmf, mf_total); + amrex::average_face_to_cellcenter( mf_avg, dcomp, srcmf, ngrow); + MultiFab::Copy( mf_avg, mf_avg, dcomp+1, dcomp+2, 1, ngrow); + MultiFab::Copy( mf_avg, *mf_total[1], 0, dcomp+1, 1, ngrow); + // Also add the real and imaginary parts of each mode. + for (int i=0 ; i < vector_field[0]->nComp() ; i++) { + MultiFab v_comp0(*vector_field[0], amrex::make_alias, i, 1); + MultiFab v_comp1(*vector_field[1], amrex::make_alias, i, 1); + MultiFab v_comp2(*vector_field[2], amrex::make_alias, i, 1); + srcmf[0] = &v_comp0; + srcmf[1] = &v_comp2; + int id = dcomp + 3*(i + 1); + amrex::average_face_to_cellcenter( mf_avg, id, srcmf, ngrow); + MultiFab::Copy( mf_avg, mf_avg, id+1, id+2, 1, ngrow); + MultiFab::Copy( mf_avg, v_comp1, 0, id+1, 1, ngrow); + } + } else { + PackPlotDataPtrs(srcmf, vector_field); + amrex::average_face_to_cellcenter( mf_avg, dcomp, srcmf, ngrow); #if (AMREX_SPACEDIM == 2) - MultiFab::Copy( mf_avg, mf_avg, dcomp+1, dcomp+2, 1, ngrow); - MultiFab::Copy( mf_avg, *vector_field[1], 0, dcomp+1, 1, ngrow); + MultiFab::Copy( mf_avg, mf_avg, dcomp+1, dcomp+2, 1, ngrow); + MultiFab::Copy( mf_avg, *vector_field[1], 0, dcomp+1, 1, ngrow); #endif + } // - Edge centered, in the same way as E on a Yee grid } else if ( !vector_field[0]->is_nodal(0) ){ - PackPlotDataPtrs(srcmf, vector_field); - amrex::average_edge_to_cellcenter( mf_avg, dcomp, srcmf, ngrow); + // See comment above, though here, the y (or theta) component + // has node centering. + if (vector_field[0]->nComp() > 1) { + // When there are more than one components, the total + // fields needs to be constructed in temporary MultiFabs + // Note that mf_total is declared in the same way as + // vector_field so that it can be passed into PackPlotDataPtrs. + std::array<std::unique_ptr<MultiFab>,3> mf_total; + mf_total[0].reset(new MultiFab(vector_field[0]->boxArray(), dm, 1, vector_field[0]->nGrowVect())); + mf_total[1].reset(new MultiFab(vector_field[1]->boxArray(), dm, 1, vector_field[1]->nGrowVect())); + mf_total[2].reset(new MultiFab(vector_field[2]->boxArray(), dm, 1, vector_field[2]->nGrowVect())); + ConstructTotalRZField(mf_total, vector_field); + PackPlotDataPtrs(srcmf, mf_total); + amrex::average_edge_to_cellcenter( mf_avg, dcomp, srcmf, ngrow); + MultiFab::Copy( mf_avg, mf_avg, dcomp+1, dcomp+2, 1, ngrow); + amrex::average_node_to_cellcenter( mf_avg, dcomp+1, + *mf_total[1], 0, 1, ngrow); + // Also add the real and imaginary parts of each mode. + for (int i=0 ; i < vector_field[0]->nComp() ; i++) { + MultiFab v_comp0(*vector_field[0], amrex::make_alias, i, 1); + MultiFab v_comp1(*vector_field[1], amrex::make_alias, i, 1); + MultiFab v_comp2(*vector_field[2], amrex::make_alias, i, 1); + srcmf[0] = &v_comp0; + srcmf[1] = &v_comp2; + int id = dcomp + 3*(i + 1); + amrex::average_edge_to_cellcenter( mf_avg, id, srcmf, ngrow); + MultiFab::Copy( mf_avg, mf_avg, id+1, id+2, 1, ngrow); + amrex::average_node_to_cellcenter( mf_avg, id+1, + v_comp1, 0, 1, ngrow); + } + } else { + PackPlotDataPtrs(srcmf, vector_field); + amrex::average_edge_to_cellcenter( mf_avg, dcomp, srcmf, ngrow); #if (AMREX_SPACEDIM == 2) - MultiFab::Copy( mf_avg, mf_avg, dcomp+1, dcomp+2, 1, ngrow); - amrex::average_node_to_cellcenter( mf_avg, dcomp+1, - *vector_field[1], 0, 1, ngrow); + MultiFab::Copy( mf_avg, mf_avg, dcomp+1, dcomp+2, 1, ngrow); + amrex::average_node_to_cellcenter( mf_avg, dcomp+1, + *vector_field[1], 0, 1, ngrow); #endif + } } else { amrex::Abort("Unknown staggering."); @@ -290,6 +371,33 @@ AverageAndPackScalarField( MultiFab& mf_avg, } } +/** \brief Add variable names to the list. + * If there are more that one mode, add the + * name of the total field and then the + * names of the real and imaginary parts of each + * mode. + */ +void +AddToVarNames( Vector<std::string>& varnames, + std::string name, std::string suffix, + int nmodes ) { + auto coords = {"x", "y", "z"}; + auto coordsRZ = {"r", "theta", "z"}; + for(auto coord:coords) varnames.push_back(name+coord+suffix); + if (nmodes > 1) { + // Note that the names are added in the same order as the fields + // are packed in AverageAndPackVectorField + for (int i = 0 ; i < nmodes ; i++) { + for(auto coord:coordsRZ) { + varnames.push_back(name + coord + suffix + std::to_string(i) + "_real"); + } + for(auto coord:coordsRZ) { + varnames.push_back(name + coord + suffix + std::to_string(i) + "_imag"); + } + } + } +} + /** \brief Write the different fields that are meant for output, * into the vector of MultiFab `mf_avg` (one MultiFab per level) * after averaging them to the cell centers. @@ -298,21 +406,29 @@ void WarpX::AverageAndPackFields ( Vector<std::string>& varnames, amrex::Vector<MultiFab>& mf_avg, const int ngrow) const { + // Factor to account for quantities that have multiple components. + // If nmodes > 1, allow space for total field and the real and + // imaginary part of each node. For now, also include the + // imaginary part of mode 0 for code symmetry, even though + // it is always zero. + int modes_factor = 1; + if (nmodes > 1) modes_factor = 2*nmodes + 1; + // Count how many different fields should be written (ncomp) const int ncomp = 0 - + static_cast<int>(plot_E_field)*3 - + static_cast<int>(plot_B_field)*3 - + static_cast<int>(plot_J_field)*3 + + static_cast<int>(plot_E_field)*3*modes_factor + + static_cast<int>(plot_B_field)*3*modes_factor + + static_cast<int>(plot_J_field)*3*modes_factor + static_cast<int>(plot_part_per_cell) + static_cast<int>(plot_part_per_grid) + static_cast<int>(plot_part_per_proc) + static_cast<int>(plot_proc_number) + static_cast<int>(plot_divb) + static_cast<int>(plot_dive) - + static_cast<int>(plot_rho) - + static_cast<int>(plot_F) - + static_cast<int>(plot_finepatch)*6 - + static_cast<int>(plot_crsepatch)*6 + + static_cast<int>(plot_rho)*modes_factor + + static_cast<int>(plot_F)*modes_factor + + static_cast<int>(plot_finepatch)*6*modes_factor + + static_cast<int>(plot_crsepatch)*6*modes_factor + static_cast<int>(costs[0] != nullptr); // Loop over levels of refinement @@ -325,19 +441,19 @@ WarpX::AverageAndPackFields ( Vector<std::string>& varnames, // add the corresponding names to `varnames` and increment dcomp int dcomp = 0; if (plot_J_field) { - AverageAndPackVectorField(mf_avg[lev], current_fp[lev], dcomp, ngrow); - if(lev==0) for(auto name:{"jx","jy","jz"}) varnames.push_back(name); - dcomp += 3; + AverageAndPackVectorField(mf_avg[lev], current_fp[lev], dmap[lev], dcomp, ngrow); + if (lev == 0) AddToVarNames(varnames, "j", "", nmodes); + dcomp += 3*modes_factor; } if (plot_E_field) { - AverageAndPackVectorField(mf_avg[lev], Efield_aux[lev], dcomp, ngrow); - if(lev==0) for(auto name:{"Ex","Ey","Ez"}) varnames.push_back(name); - dcomp += 3; + AverageAndPackVectorField(mf_avg[lev], Efield_aux[lev], dmap[lev], dcomp, ngrow); + if (lev == 0) AddToVarNames(varnames, "E", "", nmodes); + dcomp += 3*modes_factor; } if (plot_B_field) { - AverageAndPackVectorField(mf_avg[lev], Bfield_aux[lev], dcomp, ngrow); - if(lev==0) for(auto name:{"Bx","By","Bz"}) varnames.push_back(name); - dcomp += 3; + AverageAndPackVectorField(mf_avg[lev], Bfield_aux[lev], dmap[lev], dcomp, ngrow); + if (lev == 0) AddToVarNames(varnames, "B", "", nmodes); + dcomp += 3*modes_factor; } if (plot_part_per_cell) @@ -444,12 +560,12 @@ WarpX::AverageAndPackFields ( Vector<std::string>& varnames, if (plot_finepatch) { - AverageAndPackVectorField( mf_avg[lev], Efield_fp[lev], dcomp, ngrow ); - if(lev==0) for(auto name:{"Ex_fp","Ey_fp","Ez_fp"}) varnames.push_back(name); - dcomp += 3; - AverageAndPackVectorField( mf_avg[lev], Bfield_fp[lev], dcomp, ngrow ); - if(lev==0) for(auto name:{"Bx_fp","By_fp","Bz_fp"}) varnames.push_back(name); - dcomp += 3; + AverageAndPackVectorField( mf_avg[lev], Efield_fp[lev], dmap[lev], dcomp, ngrow ); + if (lev == 0) AddToVarNames(varnames, "E", "_fp", nmodes); + dcomp += 3*modes_factor; + AverageAndPackVectorField( mf_avg[lev], Bfield_fp[lev], dmap[lev], dcomp, ngrow ); + if (lev == 0) AddToVarNames(varnames, "B", "_fp", nmodes); + dcomp += 3*modes_factor; } if (plot_crsepatch) @@ -462,11 +578,11 @@ WarpX::AverageAndPackFields ( Vector<std::string>& varnames, { if (do_nodal) amrex::Abort("TODO: do_nodal && plot_crsepatch"); std::array<std::unique_ptr<MultiFab>, 3> E = getInterpolatedE(lev); - AverageAndPackVectorField( mf_avg[lev], E, dcomp, ngrow ); + AverageAndPackVectorField( mf_avg[lev], E, dmap[lev], dcomp, ngrow ); } - if(lev==0) for(auto name:{"Ex_cp","Ey_cp","Ez_cp"}) varnames.push_back(name); - dcomp += 3; + if (lev == 0) AddToVarNames(varnames, "E", "_cp", nmodes); + dcomp += 3*modes_factor; // now the magnetic field if (lev == 0) @@ -477,10 +593,10 @@ WarpX::AverageAndPackFields ( Vector<std::string>& varnames, { if (do_nodal) amrex::Abort("TODO: do_nodal && plot_crsepatch"); std::array<std::unique_ptr<MultiFab>, 3> B = getInterpolatedB(lev); - AverageAndPackVectorField( mf_avg[lev], B, dcomp, ngrow ); + AverageAndPackVectorField( mf_avg[lev], B, dmap[lev], dcomp, ngrow ); } - if(lev==0) for(auto name:{"Bx_cp","By_cp","Bz_cp"}) varnames.push_back(name); - dcomp += 3; + if (lev == 0) AddToVarNames(varnames, "B", "_cp", nmodes); + dcomp += 3*modes_factor; } if (costs[0] != nullptr) @@ -543,8 +659,8 @@ WriteRawField( const MultiFab& F, const DistributionMapping& dm, VisMF::Write(F, prefix); } else { // Copy original MultiFab into one that does not have guard cells - MultiFab tmpF( F.boxArray(), dm, 1, 0); - MultiFab::Copy(tmpF, F, 0, 0, 1, 0); + MultiFab tmpF( F.boxArray(), dm, F.nComp(), 0); + MultiFab::Copy(tmpF, F, 0, 0, F.nComp(), 0); VisMF::Write(tmpF, prefix); } @@ -566,7 +682,7 @@ WriteZeroRawField( const MultiFab& F, const DistributionMapping& dm, std::string prefix = amrex::MultiFabFileFullPrefix(lev, filename, level_prefix, field_name); - MultiFab tmpF(F.boxArray(), dm, 1, ng); + MultiFab tmpF(F.boxArray(), dm, F.nComp(), ng); tmpF.setVal(0.); VisMF::Write(tmpF, prefix); } diff --git a/Source/Diagnostics/WarpXIO.cpp b/Source/Diagnostics/WarpXIO.cpp index 24272c588..214948b2b 100644 --- a/Source/Diagnostics/WarpXIO.cpp +++ b/Source/Diagnostics/WarpXIO.cpp @@ -415,20 +415,28 @@ WarpX::GetCellCenteredData() { Vector<std::unique_ptr<MultiFab> > cc(finest_level+1); + // Factor to account for quantities that have multiple components. + // If nmodes > 1, allow space for total field and the real and + // imaginary part of each node. For now, also include the + // imaginary part of mode 0 for code symmetry, even though + // it is always zero. + int modes_factor = 1; + if (nmodes > 1) modes_factor = 2*nmodes + 1; + for (int lev = 0; lev <= finest_level; ++lev) { cc[lev].reset( new MultiFab(grids[lev], dmap[lev], nc, ng) ); int dcomp = 0; // first the electric field - AverageAndPackVectorField( *cc[lev], Efield_aux[lev], dcomp, ng ); - dcomp += 3; + AverageAndPackVectorField( *cc[lev], Efield_aux[lev], dmap[lev], dcomp, ng ); + dcomp += 3*modes_factor; // then the magnetic field - AverageAndPackVectorField( *cc[lev], Bfield_aux[lev], dcomp, ng ); - dcomp += 3; + AverageAndPackVectorField( *cc[lev], Bfield_aux[lev], dmap[lev], dcomp, ng ); + dcomp += 3*modes_factor; // then the current density - AverageAndPackVectorField( *cc[lev], current_fp[lev], dcomp, ng ); - dcomp += 3; + AverageAndPackVectorField( *cc[lev], current_fp[lev], dmap[lev], dcomp, ng ); + dcomp += 3*modes_factor; // then the charge density const std::unique_ptr<MultiFab>& charge_density = mypc->GetChargeDensity(lev); AverageAndPackScalarField( *cc[lev], *charge_density, dcomp, ng ); @@ -582,7 +590,8 @@ WarpX::WritePlotFile () const if (F_fp[lev]) WriteRawField( *F_fp[lev], dm, raw_pltname, level_prefix, "F_fp", lev, plot_raw_fields_guards); if (plot_rho) { // Use the component 1 of `rho_fp`, i.e. rho_new for time synchronization - MultiFab rho_new(*rho_fp[lev], amrex::make_alias, 1, 1); + // If nComp > 1, this is the upper half of the list of components. + MultiFab rho_new(*rho_fp[lev], amrex::make_alias, rho_fp[lev]->nComp()/2, rho_fp[lev]->nComp()/2); WriteRawField( rho_new, dm, raw_pltname, level_prefix, "rho_fp", lev, plot_raw_fields_guards); } } diff --git a/Source/Evolve/WarpXEvolveEM.cpp b/Source/Evolve/WarpXEvolveEM.cpp index 32a4747db..c5dbd73a7 100644 --- a/Source/Evolve/WarpXEvolveEM.cpp +++ b/Source/Evolve/WarpXEvolveEM.cpp @@ -348,7 +348,7 @@ WarpX::OneStep_sub1 (Real curtime) RestrictRhoFromFineToCoarsePatch(fine_lev); ApplyFilterandSumBoundaryJ(fine_lev, PatchType::fine); NodalSyncJ(fine_lev, PatchType::fine); - ApplyFilterandSumBoundaryRho(fine_lev, PatchType::fine, 0, 2); + ApplyFilterandSumBoundaryRho(fine_lev, PatchType::fine, 0, 2*nmodes); NodalSyncRho(fine_lev, PatchType::fine, 0, 2); EvolveB(fine_lev, PatchType::fine, 0.5*dt[fine_lev]); @@ -375,7 +375,7 @@ WarpX::OneStep_sub1 (Real curtime) PushParticlesandDepose(coarse_lev, curtime); StoreCurrent(coarse_lev); AddCurrentFromFineLevelandSumBoundary(coarse_lev); - AddRhoFromFineLevelandSumBoundary(coarse_lev, 0, 1); + AddRhoFromFineLevelandSumBoundary(coarse_lev, 0, nmodes); EvolveB(fine_lev, PatchType::coarse, dt[fine_lev]); EvolveF(fine_lev, PatchType::coarse, dt[fine_lev], DtType::FirstHalf); @@ -402,7 +402,7 @@ WarpX::OneStep_sub1 (Real curtime) RestrictRhoFromFineToCoarsePatch(fine_lev); ApplyFilterandSumBoundaryJ(fine_lev, PatchType::fine); NodalSyncJ(fine_lev, PatchType::fine); - ApplyFilterandSumBoundaryRho(fine_lev, PatchType::fine, 0, 2); + ApplyFilterandSumBoundaryRho(fine_lev, PatchType::fine, 0, 2*nmodes); NodalSyncRho(fine_lev, PatchType::fine, 0, 2); EvolveB(fine_lev, PatchType::fine, 0.5*dt[fine_lev]); @@ -428,7 +428,7 @@ WarpX::OneStep_sub1 (Real curtime) // by only half a coarse step (second half) RestoreCurrent(coarse_lev); AddCurrentFromFineLevelandSumBoundary(coarse_lev); - AddRhoFromFineLevelandSumBoundary(coarse_lev, 1, 1); + AddRhoFromFineLevelandSumBoundary(coarse_lev, nmodes, nmodes); EvolveE(fine_lev, PatchType::coarse, dt[fine_lev]); FillBoundaryE(fine_lev, PatchType::coarse); @@ -492,8 +492,23 @@ WarpX::ComputeDt () if (maxwell_fdtd_solver_id == 0) { // CFL time step Yee solver #ifdef WARPX_RZ - // Derived semi-analytically by R. Lehe - deltat = cfl * 1./( std::sqrt((1+0.2105)/(dx[0]*dx[0]) + 1./(dx[1]*dx[1])) * PhysConst::c ); + // In the rz case, the Courant limit has been evaluated + // semi-analytically by R. Lehe, and resulted in the following + // coefficients. For an explanation, see (not officially published) + // www.normalesup.org/~lehe/Disp_relation_Circ.pdf + // NB : Here the coefficient for m=1 as compared to this document, + // as it was observed in practice that this coefficient was not + // high enough (The simulation became unstable). + Real multimode_coeffs[6] = { 0.2105, 1.0, 3.5234, 8.5104, 15.5059, 24.5037 }; + Real multimode_alpha; + if (nmodes < 7) { + // Use the table of the coefficients + multimode_alpha = multimode_coeffs[nmodes-1]; + } else { + // Use a realistic extrapolation + multimode_alpha = (nmodes - 1)*(nmodes - 1) - 0.4; + } + deltat = cfl * 1./( std::sqrt((1+multimode_alpha)/(dx[0]*dx[0]) + 1./(dx[1]*dx[1])) * PhysConst::c ); #else deltat = cfl * 1./( std::sqrt(AMREX_D_TERM( 1./(dx[0]*dx[0]), + 1./(dx[1]*dx[1]), diff --git a/Source/FieldSolver/WarpXPushFieldsEM.cpp b/Source/FieldSolver/WarpXPushFieldsEM.cpp index c53e13f8f..fc4fb902b 100644 --- a/Source/FieldSolver/WarpXPushFieldsEM.cpp +++ b/Source/FieldSolver/WarpXPushFieldsEM.cpp @@ -109,6 +109,7 @@ WarpX::EvolveB (int lev, PatchType patch_type, amrex::Real a_dt) tbx.loVect(), tbx.hiVect(), tby.loVect(), tby.hiVect(), tbz.loVect(), tbz.hiVect(), + &nmodes, BL_TO_FORTRAN_3D((*Ex)[mfi]), BL_TO_FORTRAN_3D((*Ey)[mfi]), BL_TO_FORTRAN_3D((*Ez)[mfi]), @@ -271,6 +272,7 @@ WarpX::EvolveE (int lev, PatchType patch_type, amrex::Real a_dt) tex.loVect(), tex.hiVect(), tey.loVect(), tey.hiVect(), tez.loVect(), tez.hiVect(), + &nmodes, BL_TO_FORTRAN_3D((*Ex)[mfi]), BL_TO_FORTRAN_3D((*Ey)[mfi]), BL_TO_FORTRAN_3D((*Ez)[mfi]), diff --git a/Source/FortranInterface/WarpX_f.H b/Source/FortranInterface/WarpX_f.H index 98dd8685d..8dcd4222b 100644 --- a/Source/FortranInterface/WarpX_f.H +++ b/Source/FortranInterface/WarpX_f.H @@ -102,6 +102,7 @@ extern "C" amrex::Real* jx, const long* jx_ng, const int* jx_ntot, amrex::Real* jy, const long* jy_ng, const int* jy_ntot, amrex::Real* jz, const long* jz_ng, const int* jz_ntot, + const long* nmodes, const long* np, const amrex::Real* xp, const amrex::Real* yp, const amrex::Real* zp, const amrex::Real* uxp, const amrex::Real* uyp,const amrex::Real* uzp, @@ -117,6 +118,7 @@ extern "C" amrex::Real* jx, const long* jx_ng, const int* jx_ntot, amrex::Real* jy, const long* jy_ng, const int* jy_ntot, amrex::Real* jz, const long* jz_ng, const int* jz_ntot, + const long* nmodes, const amrex::Real* rmin, const amrex::Real* dr); @@ -136,6 +138,7 @@ extern "C" const amrex::Real* bxg, const int* bxg_lo, const int* bxg_hi, const amrex::Real* byg, const int* byg_lo, const int* byg_hi, const amrex::Real* bzg, const int* bzg_lo, const int* bzg_hi, + const long* nmodes, const int* ll4symtry, const int* l_lower_order_in_v, const int* l_nodal, const long* lvect, const long* field_gathe_algo); @@ -194,6 +197,7 @@ extern "C" const int* xlo, const int* xhi, const int* ylo, const int* yhi, const int* zlo, const int* zhi, + const long* nmodes, BL_FORT_FAB_ARG_3D(ex), BL_FORT_FAB_ARG_3D(ey), BL_FORT_FAB_ARG_3D(ez), @@ -214,6 +218,7 @@ extern "C" const int* xlo, const int* xhi, const int* ylo, const int* yhi, const int* zlo, const int* zhi, + const long* nmodes, const BL_FORT_FAB_ARG_3D(ex), const BL_FORT_FAB_ARG_3D(ey), const BL_FORT_FAB_ARG_3D(ez), diff --git a/Source/FortranInterface/WarpX_picsar.F90 b/Source/FortranInterface/WarpX_picsar.F90 index 33f85c633..e6e61b734 100644 --- a/Source/FortranInterface/WarpX_picsar.F90 +++ b/Source/FortranInterface/WarpX_picsar.F90 @@ -18,8 +18,6 @@ #define WRPX_PXR_GETEB_ENERGY_CONSERVING geteb2drz_energy_conserving_generic #define WRPX_PXR_CURRENT_DEPOSITION depose_jrjtjz_generic_rz -#define WRPX_PXR_RZ_VOLUME_SCALING_RHO apply_rz_volume_scaling_rho -#define WRPX_PXR_RZ_VOLUME_SCALING_J apply_rz_volume_scaling_j #define WRPX_PXR_PUSH_BVEC pxrpush_emrz_bvec #define WRPX_PXR_PUSH_EVEC pxrpush_emrz_evec @@ -89,6 +87,7 @@ contains ex,ey,ez,bx,by,bz,ixyzmin,xmin,ymin,zmin,dx,dy,dz,nox,noy,noz, & exg,exg_lo,exg_hi,eyg,eyg_lo,eyg_hi,ezg,ezg_lo,ezg_hi, & bxg,bxg_lo,bxg_hi,byg,byg_lo,byg_hi,bzg,bzg_lo,bzg_hi, & + nmodes, & ll4symtry,l_lower_order_in_v, l_nodal,& lvect,field_gathe_algo) & bind(C, name="warpx_geteb_energy_conserving") @@ -100,12 +99,24 @@ contains integer, intent(in) :: ixyzmin(AMREX_SPACEDIM) real(amrex_real), intent(in) :: xmin,ymin,zmin,dx,dy,dz integer(c_long), intent(in) :: field_gathe_algo - integer(c_long), intent(in) :: np,nox,noy,noz + integer(c_long), intent(in) :: np,nmodes,nox,noy,noz integer(c_int), intent(in) :: ll4symtry,l_lower_order_in_v, l_nodal integer(c_long),intent(in) :: lvect real(amrex_real), intent(in), dimension(np) :: xp,yp,zp real(amrex_real), intent(out), dimension(np) :: ex,ey,ez,bx,by,bz +#ifdef WARPX_RZ + ! The dimensions must be specified to allow the transpose + real(amrex_real),intent(in):: exg(exg_lo(1):exg_hi(1),exg_lo(2):exg_hi(2),2,nmodes) + real(amrex_real),intent(in):: eyg(eyg_lo(1):eyg_hi(1),eyg_lo(2):eyg_hi(2),2,nmodes) + real(amrex_real),intent(in):: ezg(ezg_lo(1):ezg_hi(1),ezg_lo(2):ezg_hi(2),2,nmodes) + real(amrex_real),intent(in):: bxg(bxg_lo(1):bxg_hi(1),bxg_lo(2):bxg_hi(2),2,nmodes) + real(amrex_real),intent(in):: byg(byg_lo(1):byg_hi(1),byg_lo(2):byg_hi(2),2,nmodes) + real(amrex_real),intent(in):: bzg(bzg_lo(1):bzg_hi(1),bzg_lo(2):bzg_hi(2),2,nmodes) +#else + ! These could be either 2d or 3d arrays real(amrex_real),intent(in):: exg(*), eyg(*), ezg(*), bxg(*), byg(*), bzg(*) +#endif + logical(pxr_logical) :: pxr_ll4symtry, pxr_l_lower_order_in_v, pxr_l_nodal ! Compute the number of valid cells and guard cells @@ -114,6 +125,11 @@ contains exg_nguards(AMREX_SPACEDIM), eyg_nguards(AMREX_SPACEDIM), ezg_nguards(AMREX_SPACEDIM), & bxg_nguards(AMREX_SPACEDIM), byg_nguards(AMREX_SPACEDIM), bzg_nguards(AMREX_SPACEDIM) +#ifdef WARPX_RZ + complex(amrex_real), allocatable, dimension(:,:,:) :: erg_c, etg_c, ezg_c, brg_c, btg_c, bzg_c + integer :: alloc_status +#endif + pxr_ll4symtry = ll4symtry .eq. 1 pxr_l_lower_order_in_v = l_lower_order_in_v .eq. 1 pxr_l_nodal = l_nodal .eq. 1 @@ -131,6 +147,46 @@ contains byg_nvalid = byg_lo + byg_hi - 2_c_long*ixyzmin + 1_c_long bzg_nvalid = bzg_lo + bzg_hi - 2_c_long*ixyzmin + 1_c_long +#ifdef WARPX_RZ + if (nmodes > 1) then + + allocate(erg_c(exg_lo(1):exg_hi(1),exg_lo(2):exg_hi(2),nmodes), & + etg_c(eyg_lo(1):eyg_hi(1),eyg_lo(2):eyg_hi(2),nmodes), & + ezg_c(ezg_lo(1):ezg_hi(1),ezg_lo(2):ezg_hi(2),nmodes), & + brg_c(bxg_lo(1):bxg_hi(1),bxg_lo(2):bxg_hi(2),nmodes), & + btg_c(byg_lo(1):byg_hi(1),byg_lo(2):byg_hi(2),nmodes), & + bzg_c(bzg_lo(1):bzg_hi(1),bzg_lo(2):bzg_hi(2),nmodes), stat=alloc_status) + if (alloc_status /= 0) then + print*,"Error: warpx_geteb_energy_conserving: complex arrays could not be allocated" + stop + endif + + ! Transpose the data, mapping the real and imagingary parts + ! saved separately as real numbers into complex numbers. + ! Note that the kind, amrex_real, must be specified, otherwise + ! the cmplx functions returns single precision. + erg_c(:,:,:) = cmplx(exg(:,:,1,:), exg(:,:,2,:), amrex_real) + etg_c(:,:,:) = cmplx(eyg(:,:,1,:), eyg(:,:,2,:), amrex_real) + ezg_c(:,:,:) = cmplx(ezg(:,:,1,:), ezg(:,:,2,:), amrex_real) + brg_c(:,:,:) = cmplx(bxg(:,:,1,:), bxg(:,:,2,:), amrex_real) + btg_c(:,:,:) = cmplx(byg(:,:,1,:), byg(:,:,2,:), amrex_real) + bzg_c(:,:,:) = cmplx(bzg(:,:,1,:), bzg(:,:,2,:), amrex_real) + + call geteb2dcirc_energy_conserving_generic(np, xp, yp, zp, ex, ey, ez, bx, by, bz, & + xmin, zmin, dx, dz, nmodes, nox, noz, & + pxr_l_lower_order_in_v, pxr_l_nodal, & + erg_c, exg_nguards, exg_nvalid, & + etg_c, eyg_nguards, eyg_nvalid, & + ezg_c, ezg_nguards, ezg_nvalid, & + brg_c, bxg_nguards, bxg_nvalid, & + btg_c, byg_nguards, byg_nvalid, & + bzg_c, bzg_nguards, bzg_nvalid) + + deallocate(erg_c, etg_c, ezg_c, brg_c, btg_c, bzg_c) + + else +#endif + ! Either 3d, 2d Cartesian, or purely axisymmetric CALL WRPX_PXR_GETEB_ENERGY_CONSERVING(np,xp,yp,zp, & ex,ey,ez,bx,by,bz,xmin,ymin,zmin,dx,dy,dz,nox,noy,noz, & exg,exg_nguards,exg_nvalid,& @@ -141,6 +197,9 @@ contains bzg,bzg_nguards,bzg_nvalid,& pxr_ll4symtry, pxr_l_lower_order_in_v, pxr_l_nodal, & lvect, field_gathe_algo ) +#ifdef WARPX_RZ + endif +#endif end subroutine warpx_geteb_energy_conserving @@ -271,18 +330,18 @@ subroutine warpx_charge_deposition(rho,np,xp,yp,zp,w,q,xmin,ymin,zmin,dx,dy,dz,n real(amrex_real), intent(IN) :: rmin, dr #ifdef WARPX_RZ + integer(c_long) :: type_rz_depose = 1 -#endif ! Compute the number of valid cells and guard cells integer(c_long) :: rho_nvalid(AMREX_SPACEDIM), rho_nguards(AMREX_SPACEDIM) rho_nvalid = rho_ntot - 2*rho_ng rho_nguards = rho_ng -#ifdef WARPX_RZ - CALL WRPX_PXR_RZ_VOLUME_SCALING_RHO( & + CALL apply_rz_volume_scaling_rho( & rho,rho_nguards,rho_nvalid, & rmin,dr,type_rz_depose) + #endif end subroutine warpx_charge_deposition_rz_volume_scaling @@ -313,17 +372,25 @@ subroutine warpx_charge_deposition(rho,np,xp,yp,zp,w,q,xmin,ymin,zmin,dx,dy,dz,n !> @param[in] charge_depo_algo algorithm choice for the charge deposition !> subroutine warpx_current_deposition( & - jx,jx_ng,jx_ntot,jy,jy_ng,jy_ntot,jz,jz_ng,jz_ntot, & + jx,jx_ng,jx_ntot,jy,jy_ng,jy_ntot,jz,jz_ng,jz_ntot,nmodes, & np,xp,yp,zp,uxp,uyp,uzp,gaminv,w,q,xmin,ymin,zmin,dt,dx,dy,dz,nox,noy,noz,& l_nodal,lvect,current_depo_algo) & bind(C, name="warpx_current_deposition") integer, intent(in) :: jx_ntot(AMREX_SPACEDIM), jy_ntot(AMREX_SPACEDIM), jz_ntot(AMREX_SPACEDIM) integer(c_long), intent(in) :: jx_ng, jy_ng, jz_ng + integer(c_long), intent(IN) :: nmodes integer(c_long), intent(IN) :: np integer(c_long), intent(IN) :: nox,noy,noz integer(c_int), intent(in) :: l_nodal + +#ifdef WARPX_RZ + real(amrex_real), intent(IN OUT):: jx(jx_ntot(1),jx_ntot(2),2,nmodes) + real(amrex_real), intent(IN OUT):: jy(jy_ntot(1),jy_ntot(2),2,nmodes) + real(amrex_real), intent(IN OUT):: jz(jz_ntot(1),jz_ntot(2),2,nmodes) +#else real(amrex_real), intent(IN OUT):: jx(*), jy(*), jz(*) +#endif real(amrex_real), intent(IN) :: q real(amrex_real), intent(IN) :: dx,dy,dz real(amrex_real), intent(IN) :: dt @@ -335,6 +402,13 @@ subroutine warpx_charge_deposition(rho,np,xp,yp,zp,w,q,xmin,ymin,zmin,dx,dy,dz,n integer(c_long), intent(IN) :: current_depo_algo logical(pxr_logical) :: pxr_l_nodal +#ifdef WARPX_RZ + logical(pxr_logical) :: l_particles_weight = .true. + integer(c_long) :: type_rz_depose = 1 + complex(amrex_real), allocatable, dimension(:,:,:) :: jr_c, jt_c, jz_c + integer :: alloc_status +#endif + ! Compute the number of valid cells and guard cells integer(c_long) :: jx_nvalid(AMREX_SPACEDIM), jy_nvalid(AMREX_SPACEDIM), jz_nvalid(AMREX_SPACEDIM), & jx_nguards(AMREX_SPACEDIM), jy_nguards(AMREX_SPACEDIM), jz_nguards(AMREX_SPACEDIM) @@ -357,13 +431,53 @@ subroutine warpx_charge_deposition(rho,np,xp,yp,zp,w,q,xmin,ymin,zmin,dx,dy,dz,n nox,noy,noz,pxr_l_nodal,current_depo_algo) ! Dimension 2 #elif (AMREX_SPACEDIM==2) - CALL WRPX_PXR_CURRENT_DEPOSITION( & +#ifdef WARPX_RZ + if (nmodes > 1) then + + allocate(jr_c(jx_ntot(1),jx_ntot(2),nmodes), & + jt_c(jy_ntot(1),jy_ntot(2),nmodes), & + jz_c(jz_ntot(1),jz_ntot(2),nmodes), stat=alloc_status) + if (alloc_status /= 0) then + print*,"Error: warpx_current_deposition: complex arrays could not be allocated" + stop + endif + + jr_c = 0._amrex_real + jt_c = 0._amrex_real + jz_c = 0._amrex_real + + CALL pxr_depose_jrjtjz_esirkepov_n_2d_circ( & + jr_c,jx_nguards,jx_nvalid, & + jt_c,jy_nguards,jy_nvalid, & + jz_c,jz_nguards,jz_nvalid, & + nmodes, & + np,xp,yp,zp,uxp,uyp,uzp,gaminv,w,q, & + xmin,zmin,dt,dx,dz, & + nox,noz,l_particles_weight,type_rz_depose) + + jx(:,:,1,:) = jx(:,:,1,:) + real(jr_c) + jx(:,:,2,:) = jx(:,:,2,:) + aimag(jr_c) + jy(:,:,1,:) = jy(:,:,1,:) + real(jt_c) + jy(:,:,2,:) = jy(:,:,2,:) + aimag(jt_c) + jz(:,:,1,:) = jz(:,:,1,:) + real(jz_c) + jz(:,:,2,:) = jz(:,:,2,:) + aimag(jz_c) + + deallocate(jr_c) + deallocate(jt_c) + deallocate(jz_c) + + else +#endif + CALL WRPX_PXR_CURRENT_DEPOSITION( & jx,jx_nguards,jx_nvalid, & jy,jy_nguards,jy_nvalid, & jz,jz_nguards,jz_nvalid, & np,xp,yp,zp,uxp,uyp,uzp,gaminv,w,q, & xmin,zmin,dt,dx,dz,nox,noz,pxr_l_nodal, & lvect,current_depo_algo) +#ifdef WARPX_RZ + endif +#endif #endif end subroutine warpx_current_deposition @@ -374,45 +488,92 @@ subroutine warpx_charge_deposition(rho,np,xp,yp,zp,w,q,xmin,ymin,zmin,dx,dy,dz,n !> Applies the inverse volume scaling for RZ current deposition !> !> @details - !> The scaling is done for single mode only + !> The scaling is done for all modes ! - !> @param[inout] jx,jy,jz current arrays - !> @param[in] jx_ntot,jy_ntot,jz_ntot vectors with total number of + !> @param[inout] jr,jt,jz current arrays + !> @param[in] jr_ntot,jt_ntot,jz_ntot vectors with total number of !> cells (including guard cells) along each axis for each current - !> @param[in] jx_ng,jy_ng,jz_ng vectors with number of guard cells along each + !> @param[in] jr_ng,jt_ng,jz_ng vectors with number of guard cells along each !> axis for each current !> @param[in] rmin tile grid minimum radius !> @param[in] dr radial space discretization steps !> subroutine warpx_current_deposition_rz_volume_scaling( & - jx,jx_ng,jx_ntot,jy,jy_ng,jy_ntot,jz,jz_ng,jz_ntot, & - rmin,dr) & + jr,jr_ng,jr_ntot,jt,jt_ng,jt_ntot,jz,jz_ng,jz_ntot, & + nmodes,rmin,dr) & bind(C, name="warpx_current_deposition_rz_volume_scaling") - integer, intent(in) :: jx_ntot(AMREX_SPACEDIM), jy_ntot(AMREX_SPACEDIM), jz_ntot(AMREX_SPACEDIM) - integer(c_long), intent(in) :: jx_ng, jy_ng, jz_ng - real(amrex_real), intent(IN OUT):: jx(*), jy(*), jz(*) + integer, intent(in) :: jr_ntot(AMREX_SPACEDIM), jt_ntot(AMREX_SPACEDIM), jz_ntot(AMREX_SPACEDIM) + integer(c_long), intent(in) :: jr_ng, jt_ng, jz_ng + integer(c_long), intent(in) :: nmodes + real(amrex_real), intent(IN OUT):: jr(jr_ntot(1),jr_ntot(2),2,nmodes) + real(amrex_real), intent(IN OUT):: jt(jt_ntot(1),jt_ntot(2),2,nmodes) + real(amrex_real), intent(IN OUT):: jz(jz_ntot(1),jz_ntot(2),2,nmodes) real(amrex_real), intent(IN) :: rmin, dr -#ifdef WARPX_RZ +#ifdef WARPX_RZ + + complex(amrex_real), allocatable, dimension(:,:,:) :: jr_c, jt_c, jz_c + integer :: alloc_status + integer(c_long) :: type_rz_depose = 1 -#endif + ! Compute the number of valid cells and guard cells - integer(c_long) :: jx_nvalid(AMREX_SPACEDIM), jy_nvalid(AMREX_SPACEDIM), jz_nvalid(AMREX_SPACEDIM), & - jx_nguards(AMREX_SPACEDIM), jy_nguards(AMREX_SPACEDIM), jz_nguards(AMREX_SPACEDIM) - jx_nvalid = jx_ntot - 2*jx_ng - jy_nvalid = jy_ntot - 2*jy_ng + integer(c_long) :: jr_nvalid(AMREX_SPACEDIM), jt_nvalid(AMREX_SPACEDIM), jz_nvalid(AMREX_SPACEDIM), & + jr_nguards(AMREX_SPACEDIM), jt_nguards(AMREX_SPACEDIM), jz_nguards(AMREX_SPACEDIM) + jr_nvalid = jr_ntot - 2*jr_ng + jt_nvalid = jt_ntot - 2*jt_ng jz_nvalid = jz_ntot - 2*jz_ng - jx_nguards = jx_ng - jy_nguards = jy_ng + jr_nguards = jr_ng + jt_nguards = jt_ng jz_nguards = jz_ng -#ifdef WARPX_RZ - CALL WRPX_PXR_RZ_VOLUME_SCALING_J( & - jx,jx_nguards,jx_nvalid, & - jy,jy_nguards,jy_nvalid, & + if (nmodes > 1) then + + allocate(jr_c(jr_ntot(1),jr_ntot(2),nmodes), & + jt_c(jt_ntot(1),jt_ntot(2),nmodes), & + jz_c(jz_ntot(1),jz_ntot(2),nmodes), stat=alloc_status) + if (alloc_status /= 0) then + print*,"Error: warpx_current_deposition_rz_volume_scaling: complex arrays could not be allocated" + stop + endif + + ! Transpose the data, mapping the real and imagingary parts + ! saved separately as real numbers into complex numbers. + ! Note that the kind, amrex_real, must be specified, otherwise + ! the cmplx functions returns single precision. + jr_c = cmplx(jr(:,:,1,:), jr(:,:,2,:), amrex_real) + jt_c = cmplx(jt(:,:,1,:), jt(:,:,2,:), amrex_real) + jz_c = cmplx(jz(:,:,1,:), jz(:,:,2,:), amrex_real) + + CALL apply_2dcirc_volume_scaling_j( & + jr_c, jr_nguards, jr_nvalid, & + jt_c, jt_nguards, jt_nvalid, & + jz_c, jz_nguards, jz_nvalid, & + nmodes, & + rmin,dr, & + type_rz_depose) + + ! Undo the transpose. + jr(:,:,1,:) = real(jr_c) + jr(:,:,2,:) = aimag(jr_c) + jt(:,:,1,:) = real(jt_c) + jt(:,:,2,:) = aimag(jt_c) + jz(:,:,1,:) = real(jz_c) + jz(:,:,2,:) = aimag(jz_c) + + deallocate(jr_c) + deallocate(jt_c) + deallocate(jz_c) + + else + CALL apply_rz_volume_scaling_j( & + jr,jr_nguards,jr_nvalid, & + jt,jt_nguards,jt_nvalid, & jz,jz_nguards,jz_nvalid, & rmin,dr,type_rz_depose) + endif + #endif end subroutine warpx_current_deposition_rz_volume_scaling @@ -468,6 +629,7 @@ subroutine warpx_charge_deposition(rho,np,xp,yp,zp,w,q,xmin,ymin,zmin,dx,dy,dz,n END SELECT !!!! --- push particle species positions a time step + ! Note that for RZ, the particles are pushed in 3d space #if (AMREX_SPACEDIM == 3) || (defined WARPX_RZ) CALL pxr_pushxyz(np,xp,yp,zp,uxp,uyp,uzp,gaminv,dt) #elif (AMREX_SPACEDIM == 2) @@ -548,6 +710,7 @@ subroutine warpx_charge_deposition(rho,np,xp,yp,zp,w,q,xmin,ymin,zmin,dx,dy,dz,n !> @param[in] dtsdx, dtsdy, dtsdz factors c**2 * dt/(dx, dy, dz) subroutine warpx_push_evec( & xlo, xhi, ylo, yhi, zlo, zhi, & + nmodes, & ex, exlo, exhi, & ey, eylo, eyhi, & ez, ezlo, ezhi, & @@ -567,31 +730,121 @@ subroutine warpx_charge_deposition(rho,np,xp,yp,zp,w,q,xmin,ymin,zmin,dx,dy,dz,n jxlo(BL_SPACEDIM), jxhi(BL_SPACEDIM), jylo(BL_SPACEDIM), jyhi(BL_SPACEDIM), & jzlo(BL_SPACEDIM), jzhi(BL_SPACEDIM) - real(amrex_real), intent(IN OUT):: ex(*), ey(*), ez(*) + integer(c_long), intent(in) :: nmodes - real(amrex_real), intent(IN):: bx(*), by(*), bz(*), jx(*), jy(*), jz(*) +#ifdef WARPX_RZ + ! The dimensions must be specified to allow the transpose + real(amrex_real), intent(IN OUT):: ex(exlo(1):exhi(1),exlo(2):exhi(2),2,nmodes) + real(amrex_real), intent(IN OUT):: ey(eylo(1):eyhi(1),eylo(2):eyhi(2),2,nmodes) + real(amrex_real), intent(IN OUT):: ez(ezlo(1):ezhi(1),ezlo(2):ezhi(2),2,nmodes) + real(amrex_real), intent(IN):: bx(bxlo(1):bxhi(1),bxlo(2):bxhi(2),2,nmodes) + real(amrex_real), intent(IN):: by(bylo(1):byhi(1),bylo(2):byhi(2),2,nmodes) + real(amrex_real), intent(IN):: bz(bzlo(1):bzhi(1),bzlo(2):bzhi(2),2,nmodes) + real(amrex_real), intent(IN):: jx(jxlo(1):jxhi(1),jxlo(2):jxhi(2),2,nmodes) + real(amrex_real), intent(IN):: jy(jylo(1):jyhi(1),jylo(2):jyhi(2),2,nmodes) + real(amrex_real), intent(IN):: jz(jzlo(1):jzhi(1),jzlo(2):jzhi(2),2,nmodes) +#else + ! These can be either 2d or 3d + real(amrex_real), intent(IN OUT):: ex(*) + real(amrex_real), intent(IN OUT):: ey(*) + real(amrex_real), intent(IN OUT):: ez(*) + real(amrex_real), intent(IN):: bx(*) + real(amrex_real), intent(IN):: by(*) + real(amrex_real), intent(IN):: bz(*) + real(amrex_real), intent(IN):: jx(*) + real(amrex_real), intent(IN):: jy(*) + real(amrex_real), intent(IN):: jz(*) +#endif real(amrex_real), intent(IN) :: mudt, dtsdx, dtsdy, dtsdz real(amrex_real), intent(IN) :: xmin, dx - CALL WRPX_PXR_PUSH_EVEC(& - xlo, xhi, ylo, yhi, zlo, zhi, & - ex, exlo, exhi,& - ey, eylo, eyhi,& - ez, ezlo, ezhi,& - bx, bxlo, bxhi,& - by, bylo, byhi,& - bz, bzlo, bzhi,& - jx, jxlo, jxhi,& - jy, jylo, jyhi,& - jz, jzlo, jzhi,& - mudt, dtsdx, dtsdy, dtsdz & #ifdef WARPX_RZ - ,xmin,dx & + complex(amrex_real), allocatable, dimension(:,:,:) :: er_c, et_c, ez_c + complex(amrex_real), allocatable, dimension(:,:,:) :: br_c, bt_c, bz_c + complex(amrex_real), allocatable, dimension(:,:,:) :: jr_c, jt_c, jz_c + integer :: alloc_status + + if (nmodes == 1) then +#endif + CALL WRPX_PXR_PUSH_EVEC(& + xlo, xhi, ylo, yhi, zlo, zhi, & + ex, exlo, exhi,& + ey, eylo, eyhi,& + ez, ezlo, ezhi,& + bx, bxlo, bxhi,& + by, bylo, byhi,& + bz, bzlo, bzhi,& + jx, jxlo, jxhi,& + jy, jylo, jyhi,& + jz, jzlo, jzhi,& + mudt, dtsdx, dtsdy, dtsdz & +#ifdef WARPX_RZ + ,xmin, dx & #endif - ) + ) +#ifdef WARPX_RZ + else + + allocate(er_c(exlo(1):exhi(1),exlo(2):exhi(2),nmodes), & + et_c(eylo(1):eyhi(1),eylo(2):eyhi(2),nmodes), & + ez_c(ezlo(1):ezhi(1),ezlo(2):ezhi(2),nmodes), & + br_c(bxlo(1):bxhi(1),bxlo(2):bxhi(2),nmodes), & + bt_c(bylo(1):byhi(1),bylo(2):byhi(2),nmodes), & + bz_c(bzlo(1):bzhi(1),bzlo(2):bzhi(2),nmodes), & + jr_c(jxlo(1):jxhi(1),jxlo(2):jxhi(2),nmodes), & + jt_c(jylo(1):jyhi(1),jylo(2):jyhi(2),nmodes), & + jz_c(jzlo(1):jzhi(1),jzlo(2):jzhi(2),nmodes), stat=alloc_status) + if (alloc_status /= 0) then + print*,"Error: warpx_push_evec: complex arrays could not be allocated" + stop + endif + + ! Transpose the data, mapping the real and imagingary parts + ! saved separately as real numbers into complex numbers. + ! Note that the kind, amrex_real, must be specified, otherwise + ! the cmplx functions returns single precision. + er_c = cmplx(ex(:,:,1,:), ex(:,:,2,:), amrex_real) + et_c = cmplx(ey(:,:,1,:), ey(:,:,2,:), amrex_real) + ez_c = cmplx(ez(:,:,1,:), ez(:,:,2,:), amrex_real) + br_c = cmplx(bx(:,:,1,:), bx(:,:,2,:), amrex_real) + bt_c = cmplx(by(:,:,1,:), by(:,:,2,:), amrex_real) + bz_c = cmplx(bz(:,:,1,:), bz(:,:,2,:), amrex_real) + jr_c = cmplx(jx(:,:,1,:), jx(:,:,2,:), amrex_real) + jt_c = cmplx(jy(:,:,1,:), jy(:,:,2,:), amrex_real) + jz_c = cmplx(jz(:,:,1,:), jz(:,:,2,:), amrex_real) + + CALL pxrpush_emrz_evec_multimode(& + xlo, xhi, ylo, yhi, zlo, zhi, & + nmodes, & + er_c, exlo, exhi,& + et_c, eylo, eyhi,& + ez_c, ezlo, ezhi,& + br_c, bxlo, bxhi,& + bt_c, bylo, byhi,& + bz_c, bzlo, bzhi,& + jr_c, jxlo, jxhi,& + jt_c, jylo, jyhi,& + jz_c, jzlo, jzhi,& + mudt, dtsdx, dtsdy, dtsdz, xmin, dx & + ) + + ! Only E needs to be copied back + ex(:,:,1,:) = real(er_c) + ex(:,:,2,:) = aimag(er_c) + ey(:,:,1,:) = real(et_c) + ey(:,:,2,:) = aimag(et_c) + ez(:,:,1,:) = real(ez_c) + ez(:,:,2,:) = aimag(ez_c) + + deallocate(er_c, et_c, ez_c) + deallocate(br_c, bt_c, bz_c) + deallocate(jr_c, jt_c, jz_c) + + endif +#endif end subroutine warpx_push_evec ! _________________________________________________________________ @@ -610,6 +863,7 @@ subroutine warpx_charge_deposition(rho,np,xp,yp,zp,w,q,xmin,ymin,zmin,dx,dy,dz,n !> @param[in] dtsdx, dtsdy, dtsdz factors 0.5 * dt/(dx, dy, dz) subroutine warpx_push_bvec( & xlo, xhi, ylo, yhi, zlo, zhi, & + nmodes, & ex, exlo, exhi, & ey, eylo, eyhi, & ez, ezlo, ezhi, & @@ -626,41 +880,115 @@ subroutine warpx_charge_deposition(rho,np,xp,yp,zp,w,q,xmin,ymin,zmin,dx,dy,dz,n bylo(BL_SPACEDIM), byhi(BL_SPACEDIM), bzlo(BL_SPACEDIM), bzhi(BL_SPACEDIM), & maxwell_fdtd_solver_id - real(amrex_real), intent(IN OUT):: ex(*), ey(*), ez(*) + integer(c_long), intent(in) :: nmodes - real(amrex_real), intent(IN):: bx(*), by(*), bz(*) +#ifdef WARPX_RZ + ! The dimensions must be specified to allow the transpose + real(amrex_real), intent(IN):: ex(exlo(1):exhi(1),exlo(2):exhi(2),2,nmodes) + real(amrex_real), intent(IN):: ey(eylo(1):eyhi(1),eylo(2):eyhi(2),2,nmodes) + real(amrex_real), intent(IN):: ez(ezlo(1):ezhi(1),ezlo(2):ezhi(2),2,nmodes) + real(amrex_real), intent(IN OUT):: bx(bxlo(1):bxhi(1),bxlo(2):bxhi(2),2,nmodes) + real(amrex_real), intent(IN OUT):: by(bylo(1):byhi(1),bylo(2):byhi(2),2,nmodes) + real(amrex_real), intent(IN OUT):: bz(bzlo(1):bzhi(1),bzlo(2):bzhi(2),2,nmodes) +#else + real(amrex_real), intent(IN):: ex(*) + real(amrex_real), intent(IN):: ey(*) + real(amrex_real), intent(IN):: ez(*) + real(amrex_real), intent(IN OUT):: bx(*) + real(amrex_real), intent(IN OUT):: by(*) + real(amrex_real), intent(IN OUT):: bz(*) +#endif real(amrex_real), intent(IN) :: dtsdx, dtsdy, dtsdz real(amrex_real), intent(IN) :: xmin, dx - IF (maxwell_fdtd_solver_id .eq. 0) THEN - ! Yee FDTD solver - CALL WRPX_PXR_PUSH_BVEC( & - xlo, xhi, ylo, yhi, zlo, zhi, & - ex, exlo, exhi, & - ey, eylo, eyhi, & - ez, ezlo, ezhi, & - bx, bxlo, bxhi, & - by, bylo, byhi, & - bz, bzlo, bzhi, & - dtsdx,dtsdy,dtsdz & #ifdef WARPX_RZ - ,xmin,dx & + complex(amrex_real), allocatable, dimension(:,:,:) :: er_c, et_c, ez_c + complex(amrex_real), allocatable, dimension(:,:,:) :: br_c, bt_c, bz_c + integer :: alloc_status #endif - ) - ELSE IF (maxwell_fdtd_solver_id .eq. 1) THEN - ! Cole-Karkkainen FDTD solver - CALL WRPX_PXR_PUSH_BVEC_CKC( & - xlo, xhi, ylo, yhi, zlo, zhi, & - ex, exlo, exhi, & - ey, eylo, eyhi, & - ez, ezlo, ezhi, & - bx, bxlo, bxhi, & - by, bylo, byhi, & - bz, bzlo, bzhi, & - dtsdx,dtsdy,dtsdz) - ENDIF + + if (nmodes == 1) then + + IF (maxwell_fdtd_solver_id .eq. 0) THEN + ! Yee FDTD solver + CALL WRPX_PXR_PUSH_BVEC( & + xlo, xhi, ylo, yhi, zlo, zhi, & + ex, exlo, exhi, & + ey, eylo, eyhi, & + ez, ezlo, ezhi, & + bx, bxlo, bxhi, & + by, bylo, byhi, & + bz, bzlo, bzhi, & + dtsdx,dtsdy,dtsdz & +#ifdef WARPX_RZ + ,xmin,dx & +#endif + ) + ELSE IF (maxwell_fdtd_solver_id .eq. 1) THEN + ! Cole-Karkkainen FDTD solver + CALL WRPX_PXR_PUSH_BVEC_CKC( & + xlo, xhi, ylo, yhi, zlo, zhi, & + ex, exlo, exhi, & + ey, eylo, eyhi, & + ez, ezlo, ezhi, & + bx, bxlo, bxhi, & + by, bylo, byhi, & + bz, bzlo, bzhi, & + dtsdx,dtsdy,dtsdz) + ENDIF + +#ifdef WARPX_RZ + else + + allocate(er_c(exlo(1):exhi(1),exlo(2):exhi(2),nmodes), & + et_c(eylo(1):eyhi(1),eylo(2):eyhi(2),nmodes), & + ez_c(ezlo(1):ezhi(1),ezlo(2):ezhi(2),nmodes), & + br_c(bxlo(1):bxhi(1),bxlo(2):bxhi(2),nmodes), & + bt_c(bylo(1):byhi(1),bylo(2):byhi(2),nmodes), & + bz_c(bzlo(1):bzhi(1),bzlo(2):bzhi(2),nmodes), stat=alloc_status) + if (alloc_status /= 0) then + print*,"Error: warpx_push_bvec: complex arrays could not be allocated" + stop + endif + + ! Transpose the data, mapping the real and imagingary parts + ! saved separately as real numbers into complex numbers. + ! Note that the kind, amrex_real, must be specified, otherwise + ! the cmplx functions returns single precision. + er_c = cmplx(ex(:,:,1,:), ex(:,:,2,:), amrex_real) + et_c = cmplx(ey(:,:,1,:), ey(:,:,2,:), amrex_real) + ez_c = cmplx(ez(:,:,1,:), ez(:,:,2,:), amrex_real) + br_c = cmplx(bx(:,:,1,:), bx(:,:,2,:), amrex_real) + bt_c = cmplx(by(:,:,1,:), by(:,:,2,:), amrex_real) + bz_c = cmplx(bz(:,:,1,:), bz(:,:,2,:), amrex_real) + + CALL pxrpush_emrz_bvec_multimode(& + xlo, xhi, ylo, yhi, zlo, zhi, & + nmodes, & + er_c, exlo, exhi,& + et_c, eylo, eyhi,& + ez_c, ezlo, ezhi,& + br_c, bxlo, bxhi,& + bt_c, bylo, byhi,& + bz_c, bzlo, bzhi,& + dtsdx, dtsdy, dtsdz, xmin, dx & + ) + + ! Only B needs to be copied back + bx(:,:,1,:) = real(br_c) + bx(:,:,2,:) = aimag(br_c) + by(:,:,1,:) = real(bt_c) + by(:,:,2,:) = aimag(bt_c) + bz(:,:,1,:) = real(bz_c) + bz(:,:,2,:) = aimag(bz_c) + + deallocate(er_c, et_c, ez_c) + deallocate(br_c, bt_c, bz_c) + +#endif + endif end subroutine warpx_push_bvec ! _________________________________________________________________ diff --git a/Source/Initialization/PlasmaInjector.cpp b/Source/Initialization/PlasmaInjector.cpp index f9642d1b6..a7dbe728a 100644 --- a/Source/Initialization/PlasmaInjector.cpp +++ b/Source/Initialization/PlasmaInjector.cpp @@ -218,7 +218,7 @@ void RegularPosition::getPositionUnitBox(vec3& r, int i_part, int ref_fac) { int nx = ref_fac*_num_particles_per_cell_each_dim[0]; int ny = ref_fac*_num_particles_per_cell_each_dim[1]; -#if AMREX_SPACEDIM == 3 +#if (AMREX_SPACEDIM == 3) || (defined WARPX_RZ) int nz = ref_fac*_num_particles_per_cell_each_dim[2]; #else int nz = 1; @@ -296,9 +296,9 @@ PlasmaInjector::PlasmaInjector(int ispecies, const std::string& name) parseDensity(pp); parseMomentum(pp); } else if (part_pos_s == "nuniformpercell") { - num_particles_per_cell_each_dim.resize(3); + num_particles_per_cell_each_dim.assign(3, 1); pp.getarr("num_particles_per_cell_each_dim", num_particles_per_cell_each_dim); -#if ( AMREX_SPACEDIM == 2 ) +#if ( AMREX_SPACEDIM == 2 ) && !defined(WARPX_RZ) num_particles_per_cell_each_dim[2] = 1; #endif part_pos.reset(new RegularPosition(num_particles_per_cell_each_dim)); diff --git a/Source/Parallelization/WarpXComm.cpp b/Source/Parallelization/WarpXComm.cpp index 9d85783b0..0ca1e8a5d 100644 --- a/Source/Parallelization/WarpXComm.cpp +++ b/Source/Parallelization/WarpXComm.cpp @@ -59,24 +59,24 @@ WarpX::UpdateAuxilaryData () // B field { - MultiFab dBx(Bfield_cp[lev][0]->boxArray(), dm, 1, ng); - MultiFab dBy(Bfield_cp[lev][1]->boxArray(), dm, 1, ng); - MultiFab dBz(Bfield_cp[lev][2]->boxArray(), dm, 1, ng); + MultiFab dBx(Bfield_cp[lev][0]->boxArray(), dm, Bfield_cp[lev][0]->nComp(), ng); + MultiFab dBy(Bfield_cp[lev][1]->boxArray(), dm, Bfield_cp[lev][1]->nComp(), ng); + MultiFab dBz(Bfield_cp[lev][2]->boxArray(), dm, Bfield_cp[lev][2]->nComp(), ng); dBx.setVal(0.0); dBy.setVal(0.0); dBz.setVal(0.0); - dBx.ParallelCopy(*Bfield_aux[lev-1][0], 0, 0, 1, ng, ng, crse_period); - dBy.ParallelCopy(*Bfield_aux[lev-1][1], 0, 0, 1, ng, ng, crse_period); - dBz.ParallelCopy(*Bfield_aux[lev-1][2], 0, 0, 1, ng, ng, crse_period); + dBx.ParallelCopy(*Bfield_aux[lev-1][0], 0, 0, Bfield_aux[lev-1][0]->nComp(), ng, ng, crse_period); + dBy.ParallelCopy(*Bfield_aux[lev-1][1], 0, 0, Bfield_aux[lev-1][1]->nComp(), ng, ng, crse_period); + dBz.ParallelCopy(*Bfield_aux[lev-1][2], 0, 0, Bfield_aux[lev-1][2]->nComp(), ng, ng, crse_period); if (Bfield_cax[lev][0]) { - MultiFab::Copy(*Bfield_cax[lev][0], dBx, 0, 0, 1, ng); - MultiFab::Copy(*Bfield_cax[lev][1], dBy, 0, 0, 1, ng); - MultiFab::Copy(*Bfield_cax[lev][2], dBz, 0, 0, 1, ng); + MultiFab::Copy(*Bfield_cax[lev][0], dBx, 0, 0, Bfield_cax[lev][0]->nComp(), ng); + MultiFab::Copy(*Bfield_cax[lev][1], dBy, 0, 0, Bfield_cax[lev][1]->nComp(), ng); + MultiFab::Copy(*Bfield_cax[lev][2], dBz, 0, 0, Bfield_cax[lev][2]->nComp(), ng); } - MultiFab::Subtract(dBx, *Bfield_cp[lev][0], 0, 0, 1, ng); - MultiFab::Subtract(dBy, *Bfield_cp[lev][1], 0, 0, 1, ng); - MultiFab::Subtract(dBz, *Bfield_cp[lev][2], 0, 0, 1, ng); + MultiFab::Subtract(dBx, *Bfield_cp[lev][0], 0, 0, Bfield_cp[lev][0]->nComp(), ng); + MultiFab::Subtract(dBy, *Bfield_cp[lev][1], 0, 0, Bfield_cp[lev][1]->nComp(), ng); + MultiFab::Subtract(dBz, *Bfield_cp[lev][2], 0, 0, Bfield_cp[lev][2]->nComp(), ng); const Real* dx = Geom(lev-1).CellSize(); const int refinement_ratio = refRatio(lev-1)[0]; @@ -134,24 +134,24 @@ WarpX::UpdateAuxilaryData () // E field { - MultiFab dEx(Efield_cp[lev][0]->boxArray(), dm, 1, ng); - MultiFab dEy(Efield_cp[lev][1]->boxArray(), dm, 1, ng); - MultiFab dEz(Efield_cp[lev][2]->boxArray(), dm, 1, ng); + MultiFab dEx(Efield_cp[lev][0]->boxArray(), dm, Efield_cp[lev][0]->nComp(), ng); + MultiFab dEy(Efield_cp[lev][1]->boxArray(), dm, Efield_cp[lev][1]->nComp(), ng); + MultiFab dEz(Efield_cp[lev][2]->boxArray(), dm, Efield_cp[lev][2]->nComp(), ng); dEx.setVal(0.0); dEy.setVal(0.0); dEz.setVal(0.0); - dEx.ParallelCopy(*Efield_aux[lev-1][0], 0, 0, 1, ng, ng, crse_period); - dEy.ParallelCopy(*Efield_aux[lev-1][1], 0, 0, 1, ng, ng, crse_period); - dEz.ParallelCopy(*Efield_aux[lev-1][2], 0, 0, 1, ng, ng, crse_period); + dEx.ParallelCopy(*Efield_aux[lev-1][0], 0, 0, Efield_aux[lev-1][0]->nComp(), ng, ng, crse_period); + dEy.ParallelCopy(*Efield_aux[lev-1][1], 0, 0, Efield_aux[lev-1][1]->nComp(), ng, ng, crse_period); + dEz.ParallelCopy(*Efield_aux[lev-1][2], 0, 0, Efield_aux[lev-1][2]->nComp(), ng, ng, crse_period); if (Efield_cax[lev][0]) { - MultiFab::Copy(*Efield_cax[lev][0], dEx, 0, 0, 1, ng); - MultiFab::Copy(*Efield_cax[lev][1], dEy, 0, 0, 1, ng); - MultiFab::Copy(*Efield_cax[lev][2], dEz, 0, 0, 1, ng); + MultiFab::Copy(*Efield_cax[lev][0], dEx, 0, 0, Efield_cax[lev][0]->nComp(), ng); + MultiFab::Copy(*Efield_cax[lev][1], dEy, 0, 0, Efield_cax[lev][1]->nComp(), ng); + MultiFab::Copy(*Efield_cax[lev][2], dEz, 0, 0, Efield_cax[lev][2]->nComp(), ng); } - MultiFab::Subtract(dEx, *Efield_cp[lev][0], 0, 0, 1, ng); - MultiFab::Subtract(dEy, *Efield_cp[lev][1], 0, 0, 1, ng); - MultiFab::Subtract(dEz, *Efield_cp[lev][2], 0, 0, 1, ng); + MultiFab::Subtract(dEx, *Efield_cp[lev][0], 0, 0, Efield_cp[lev][0]->nComp(), ng); + MultiFab::Subtract(dEy, *Efield_cp[lev][1], 0, 0, Efield_cp[lev][1]->nComp(), ng); + MultiFab::Subtract(dEz, *Efield_cp[lev][2], 0, 0, Efield_cp[lev][2]->nComp(), ng); const int refinement_ratio = refRatio(lev-1)[0]; #ifdef _OPEMP @@ -199,8 +199,8 @@ WarpX::UpdateAuxilaryData () FArrayBox& aux = (*Efield_aux[lev][idim])[mfi]; FArrayBox& fp = (*Efield_fp[lev][idim])[mfi]; const Box& bx = aux.box(); - aux.copy(fp, bx, 0, bx, 0, 1); - aux.plus(efab[idim], bx, bx, 0, 0, 1); + aux.copy(fp, bx, 0, bx, 0, Efield_fp[lev][idim]->nComp()); + aux.plus(efab[idim], bx, bx, 0, 0, Efield_fp[lev][idim]->nComp()); } } } @@ -409,7 +409,7 @@ WarpX::SyncCurrent (const std::array<const amrex::MultiFab*,3>& fine, ffab.resize(fbx); fbx &= (*fine[idim])[mfi].box(); ffab.setVal(0.0); - ffab.copy((*fine[idim])[mfi], fbx, 0, fbx, 0, 1); + ffab.copy((*fine[idim])[mfi], fbx, 0, fbx, 0, fine[idim]->nComp()); WRPX_SYNC_CURRENT(bx.loVect(), bx.hiVect(), BL_TO_FORTRAN_ANYD((*crse[idim])[mfi]), BL_TO_FORTRAN_ANYD(ffab), @@ -505,11 +505,11 @@ WarpX::ApplyFilterandSumBoundaryJ (int lev, PatchType patch_type) if (use_filter) { IntVect ng = j[idim]->nGrowVect(); ng += bilinear_filter.stencil_length_each_dir-1; - MultiFab jf(j[idim]->boxArray(), j[idim]->DistributionMap(), 1, ng); + MultiFab jf(j[idim]->boxArray(), j[idim]->DistributionMap(), j[idim]->nComp(), ng); bilinear_filter.ApplyStencil(jf, *j[idim]); - WarpXSumGuardCells(*(j[idim]), jf, period); + WarpXSumGuardCells(*(j[idim]), jf, period, 0, (j[idim])->nComp()); } else { - WarpXSumGuardCells(*(j[idim]), period); + WarpXSumGuardCells(*(j[idim]), period, 0, (j[idim])->nComp()); } } } @@ -539,7 +539,7 @@ WarpX::AddCurrentFromFineLevelandSumBoundary (int lev) const auto& period = Geom(lev).periodicity(); for (int idim = 0; idim < 3; ++idim) { MultiFab mf(current_fp[lev][idim]->boxArray(), - current_fp[lev][idim]->DistributionMap(), 1, 0); + current_fp[lev][idim]->DistributionMap(), current_fp[lev][idim]->nComp(), 0); mf.setVal(0.0); if (use_filter && current_buf[lev+1][idim]) { @@ -547,18 +547,18 @@ WarpX::AddCurrentFromFineLevelandSumBoundary (int lev) IntVect ng = current_cp[lev+1][idim]->nGrowVect(); ng += bilinear_filter.stencil_length_each_dir-1; MultiFab jfc(current_cp[lev+1][idim]->boxArray(), - current_cp[lev+1][idim]->DistributionMap(), 1, ng); + current_cp[lev+1][idim]->DistributionMap(), current_cp[lev+1][idim]->nComp(), ng); bilinear_filter.ApplyStencil(jfc, *current_cp[lev+1][idim]); // buffer patch of fine level MultiFab jfb(current_buf[lev+1][idim]->boxArray(), - current_buf[lev+1][idim]->DistributionMap(), 1, ng); + current_buf[lev+1][idim]->DistributionMap(), current_buf[lev+1][idim]->nComp(), ng); bilinear_filter.ApplyStencil(jfb, *current_buf[lev+1][idim]); - MultiFab::Add(jfb, jfc, 0, 0, 1, ng); - mf.ParallelAdd(jfb, 0, 0, 1, ng, IntVect::TheZeroVector(), period); + MultiFab::Add(jfb, jfc, 0, 0, current_buf[lev+1][idim]->nComp(), ng); + mf.ParallelAdd(jfb, 0, 0, current_buf[lev+1][idim]->nComp(), ng, IntVect::TheZeroVector(), period); - WarpXSumGuardCells(*current_cp[lev+1][idim], jfc, period); + WarpXSumGuardCells(*current_cp[lev+1][idim], jfc, period, 0, current_cp[lev+1][idim]->nComp()); } else if (use_filter) // but no buffer { @@ -566,29 +566,29 @@ WarpX::AddCurrentFromFineLevelandSumBoundary (int lev) IntVect ng = current_cp[lev+1][idim]->nGrowVect(); ng += bilinear_filter.stencil_length_each_dir-1; MultiFab jf(current_cp[lev+1][idim]->boxArray(), - current_cp[lev+1][idim]->DistributionMap(), 1, ng); + current_cp[lev+1][idim]->DistributionMap(), current_cp[lev+1][idim]->nComp(), ng); bilinear_filter.ApplyStencil(jf, *current_cp[lev+1][idim]); - mf.ParallelAdd(jf, 0, 0, 1, ng, IntVect::TheZeroVector(), period); - WarpXSumGuardCells(*current_cp[lev+1][idim], jf, period); + mf.ParallelAdd(jf, 0, 0, current_cp[lev+1][idim]->nComp(), ng, IntVect::TheZeroVector(), period); + WarpXSumGuardCells(*current_cp[lev+1][idim], jf, period, 0, current_cp[lev+1][idim]->nComp()); } else if (current_buf[lev+1][idim]) // but no filter { MultiFab::Copy(*current_buf[lev+1][idim], - *current_cp [lev+1][idim], 0, 0, 1, + *current_cp [lev+1][idim], 0, 0, current_buf[lev+1][idim]->nComp(), current_cp[lev+1][idim]->nGrow()); - mf.ParallelAdd(*current_buf[lev+1][idim], 0, 0, 1, + mf.ParallelAdd(*current_buf[lev+1][idim], 0, 0, current_buf[lev+1][idim]->nComp(), current_buf[lev+1][idim]->nGrowVect(), IntVect::TheZeroVector(), period); - WarpXSumGuardCells(*(current_cp[lev+1][idim]), period); + WarpXSumGuardCells(*(current_cp[lev+1][idim]), period, 0, current_cp[lev+1][idim]->nComp()); } else // no filter, no buffer { - mf.ParallelAdd(*current_cp[lev+1][idim], 0, 0, 1, + mf.ParallelAdd(*current_cp[lev+1][idim], 0, 0, current_cp[lev+1][idim]->nComp(), current_cp[lev+1][idim]->nGrowVect(), IntVect::TheZeroVector(), period); - WarpXSumGuardCells(*(current_cp[lev+1][idim]), period); + WarpXSumGuardCells(*(current_cp[lev+1][idim]), period, 0, current_cp[lev+1][idim]->nComp()); } - MultiFab::Add(*current_fp[lev][idim], mf, 0, 0, 1, 0); + MultiFab::Add(*current_fp[lev][idim], mf, 0, 0, current_fp[lev+1][idim]->nComp(), 0); } NodalSyncJ(lev+1, PatchType::coarse); } diff --git a/Source/Parallelization/WarpXRegrid.cpp b/Source/Parallelization/WarpXRegrid.cpp index 8d7873041..eb119d4a2 100644 --- a/Source/Parallelization/WarpXRegrid.cpp +++ b/Source/Parallelization/WarpXRegrid.cpp @@ -46,21 +46,21 @@ WarpX::RemakeLevel (int lev, Real time, const BoxArray& ba, const DistributionMa { const IntVect& ng = Bfield_fp[lev][idim]->nGrowVect(); auto pmf = std::unique_ptr<MultiFab>(new MultiFab(Bfield_fp[lev][idim]->boxArray(), - dm, 1, ng)); - pmf->Redistribute(*Bfield_fp[lev][idim], 0, 0, 1, ng); + dm, Bfield_fp[lev][idim]->nComp(), ng)); + pmf->Redistribute(*Bfield_fp[lev][idim], 0, 0, Bfield_fp[lev][idim]->nComp(), ng); Bfield_fp[lev][idim] = std::move(pmf); } { const IntVect& ng = Efield_fp[lev][idim]->nGrowVect(); auto pmf = std::unique_ptr<MultiFab>(new MultiFab(Efield_fp[lev][idim]->boxArray(), - dm, 1, ng)); - pmf->Redistribute(*Efield_fp[lev][idim], 0, 0, 1, ng); + dm, Efield_fp[lev][idim]->nComp(), ng)); + pmf->Redistribute(*Efield_fp[lev][idim], 0, 0, Efield_fp[lev][idim]->nComp(), ng); Efield_fp[lev][idim] = std::move(pmf); } { const IntVect& ng = current_fp[lev][idim]->nGrowVect(); auto pmf = std::unique_ptr<MultiFab>(new MultiFab(current_fp[lev][idim]->boxArray(), - dm, 1, ng)); + dm, current_fp[lev][idim]->nComp(), ng)); current_fp[lev][idim] = std::move(pmf); current_fp_owner_masks[lev][idim] = std::move(current_fp[lev][idim]->OwnerMask(period)); } @@ -68,7 +68,7 @@ WarpX::RemakeLevel (int lev, Real time, const BoxArray& ba, const DistributionMa { const IntVect& ng = current_store[lev][idim]->nGrowVect(); auto pmf = std::unique_ptr<MultiFab>(new MultiFab(current_store[lev][idim]->boxArray(), - dm, 1, ng)); + dm, current_store[lev][idim]->nComp(), ng)); // no need to redistribute current_store[lev][idim] = std::move(pmf); } @@ -77,8 +77,8 @@ WarpX::RemakeLevel (int lev, Real time, const BoxArray& ba, const DistributionMa if (F_fp[lev] != nullptr) { const IntVect& ng = F_fp[lev]->nGrowVect(); auto pmf = std::unique_ptr<MultiFab>(new MultiFab(F_fp[lev]->boxArray(), - dm, 1, ng)); - pmf->Redistribute(*F_fp[lev], 0, 0, 1, ng); + dm, F_fp[lev]->nComp(), ng)); + pmf->Redistribute(*F_fp[lev], 0, 0, F_fp[lev]->nComp(), ng); F_fp[lev] = std::move(pmf); } @@ -96,8 +96,8 @@ WarpX::RemakeLevel (int lev, Real time, const BoxArray& ba, const DistributionMa if (lev == 0) { for (int idim = 0; idim < 3; ++idim) { - Bfield_aux[lev][idim].reset(new MultiFab(*Bfield_fp[lev][idim], amrex::make_alias, 0, 1)); - Efield_aux[lev][idim].reset(new MultiFab(*Efield_fp[lev][idim], amrex::make_alias, 0, 1)); + Bfield_aux[lev][idim].reset(new MultiFab(*Bfield_fp[lev][idim], amrex::make_alias, 0, Bfield_aux[lev][idim]->nComp())); + Efield_aux[lev][idim].reset(new MultiFab(*Efield_fp[lev][idim], amrex::make_alias, 0, Efield_aux[lev][idim]->nComp())); } } else { for (int idim=0; idim < 3; ++idim) @@ -105,15 +105,15 @@ WarpX::RemakeLevel (int lev, Real time, const BoxArray& ba, const DistributionMa { const IntVect& ng = Bfield_aux[lev][idim]->nGrowVect(); auto pmf = std::unique_ptr<MultiFab>(new MultiFab(Bfield_aux[lev][idim]->boxArray(), - dm, 1, ng)); - // pmf->Redistribute(*Bfield_aux[lev][idim], 0, 0, 1, ng); + dm, Bfield_aux[lev][idim]->nComp(), ng)); + // pmf->Redistribute(*Bfield_aux[lev][idim], 0, 0, Bfield_aux[lev][idim]->nComp(), ng); Bfield_aux[lev][idim] = std::move(pmf); } { const IntVect& ng = Efield_aux[lev][idim]->nGrowVect(); auto pmf = std::unique_ptr<MultiFab>(new MultiFab(Efield_aux[lev][idim]->boxArray(), - dm, 1, ng)); - // pmf->Redistribute(*Efield_aux[lev][idim], 0, 0, 1, ng); + dm, Efield_aux[lev][idim]->nComp(), ng)); + // pmf->Redistribute(*Efield_aux[lev][idim], 0, 0, Efield_aux[lev][idim]->nComp(), ng); Efield_aux[lev][idim] = std::move(pmf); } } @@ -127,21 +127,21 @@ WarpX::RemakeLevel (int lev, Real time, const BoxArray& ba, const DistributionMa { const IntVect& ng = Bfield_cp[lev][idim]->nGrowVect(); auto pmf = std::unique_ptr<MultiFab>(new MultiFab(Bfield_cp[lev][idim]->boxArray(), - dm, 1, ng)); - pmf->Redistribute(*Bfield_cp[lev][idim], 0, 0, 1, ng); + dm, Bfield_cp[lev][idim]->nComp(), ng)); + pmf->Redistribute(*Bfield_cp[lev][idim], 0, 0, Bfield_cp[lev][idim]->nComp(), ng); Bfield_cp[lev][idim] = std::move(pmf); } { const IntVect& ng = Efield_cp[lev][idim]->nGrowVect(); auto pmf = std::unique_ptr<MultiFab>(new MultiFab(Efield_cp[lev][idim]->boxArray(), - dm, 1, ng)); - pmf->Redistribute(*Efield_cp[lev][idim], 0, 0, 1, ng); + dm, Efield_cp[lev][idim]->nComp(), ng)); + pmf->Redistribute(*Efield_cp[lev][idim], 0, 0, Efield_cp[lev][idim]->nComp(), ng); Efield_cp[lev][idim] = std::move(pmf); } { const IntVect& ng = current_cp[lev][idim]->nGrowVect(); auto pmf = std::unique_ptr<MultiFab>( new MultiFab(current_cp[lev][idim]->boxArray(), - dm, 1, ng)); + dm, current_cp[lev][idim]->nComp(), ng)); current_cp[lev][idim] = std::move(pmf); current_cp_owner_masks[lev][idim] = std::move( current_cp[lev][idim]->OwnerMask(cperiod)); @@ -151,8 +151,8 @@ WarpX::RemakeLevel (int lev, Real time, const BoxArray& ba, const DistributionMa if (F_cp[lev] != nullptr) { const IntVect& ng = F_cp[lev]->nGrowVect(); auto pmf = std::unique_ptr<MultiFab>(new MultiFab(F_cp[lev]->boxArray(), - dm, 1, ng)); - pmf->Redistribute(*F_cp[lev], 0, 0, 1, ng); + dm, F_cp[lev]->nComp(), ng)); + pmf->Redistribute(*F_cp[lev], 0, 0, F_cp[lev]->nComp(), ng); F_cp[lev] = std::move(pmf); } @@ -173,24 +173,24 @@ WarpX::RemakeLevel (int lev, Real time, const BoxArray& ba, const DistributionMa { const IntVect& ng = Bfield_cax[lev][idim]->nGrowVect(); auto pmf = std::unique_ptr<MultiFab>(new MultiFab(Bfield_cax[lev][idim]->boxArray(), - dm, 1, ng)); - // pmf->ParallelCopy(*Bfield_cax[lev][idim], 0, 0, 1, ng, ng); + dm, Bfield_cax[lev][idim]->nComp(), ng)); + // pmf->ParallelCopy(*Bfield_cax[lev][idim], 0, 0, Bfield_cax[lev][idim]->nComp(), ng, ng); Bfield_cax[lev][idim] = std::move(pmf); } if (Efield_cax[lev][idim]) { const IntVect& ng = Efield_cax[lev][idim]->nGrowVect(); auto pmf = std::unique_ptr<MultiFab>(new MultiFab(Efield_cax[lev][idim]->boxArray(), - dm, 1, ng)); - // pmf->ParallelCopy(*Efield_cax[lev][idim], 0, 0, 1, ng, ng); + dm, Efield_cax[lev][idim]->nComp(), ng)); + // pmf->ParallelCopy(*Efield_cax[lev][idim], 0, 0, Efield_cax[lev][idim]->nComp(), ng, ng); Efield_cax[lev][idim] = std::move(pmf); } if (current_buf[lev][idim]) { const IntVect& ng = current_buf[lev][idim]->nGrowVect(); auto pmf = std::unique_ptr<MultiFab>(new MultiFab(current_buf[lev][idim]->boxArray(), - dm, 1, ng)); - // pmf->ParallelCopy(*current_buf[lev][idim], 0, 0, 1, ng, ng); + dm, current_buf[lev][idim]->nComp(), ng)); + // pmf->ParallelCopy(*current_buf[lev][idim], 0, 0, current_buf[lev][idim]->nComp(), ng, ng); current_buf[lev][idim] = std::move(pmf); } } @@ -198,24 +198,24 @@ WarpX::RemakeLevel (int lev, Real time, const BoxArray& ba, const DistributionMa { const IntVect& ng = charge_buf[lev]->nGrowVect(); auto pmf = std::unique_ptr<MultiFab>(new MultiFab(charge_buf[lev]->boxArray(), - dm, 1, ng)); - // pmf->ParallelCopy(*charge_buf[lev][idim], 0, 0, 1, ng, ng); + dm, charge_buf[lev]->nComp(), ng)); + // pmf->ParallelCopy(*charge_buf[lev][idim], 0, 0, charge_buf[lev]->nComp(), ng, ng); charge_buf[lev] = std::move(pmf); } if (current_buffer_masks[lev]) { const IntVect& ng = current_buffer_masks[lev]->nGrowVect(); auto pmf = std::unique_ptr<iMultiFab>(new iMultiFab(current_buffer_masks[lev]->boxArray(), - dm, 1, ng)); - // pmf->ParallelCopy(*current_buffer_masks[lev], 0, 0, 1, ng, ng); + dm, current_buffer_masks[lev]->nComp(), ng)); + // pmf->ParallelCopy(*current_buffer_masks[lev], 0, 0, current_buffer_masks[lev]->nComp(), ng, ng); current_buffer_masks[lev] = std::move(pmf); } if (gather_buffer_masks[lev]) { const IntVect& ng = gather_buffer_masks[lev]->nGrowVect(); auto pmf = std::unique_ptr<iMultiFab>(new iMultiFab(gather_buffer_masks[lev]->boxArray(), - dm, 1, ng)); - // pmf->ParallelCopy(*gather_buffer_masks[lev], 0, 0, 1, ng, ng); + dm, gather_buffer_masks[lev]->nComp(), ng)); + // pmf->ParallelCopy(*gather_buffer_masks[lev], 0, 0, gather_buffer_masks[lev]->nComp(), ng, ng); gather_buffer_masks[lev] = std::move(pmf); } } diff --git a/Source/Particles/MultiParticleContainer.cpp b/Source/Particles/MultiParticleContainer.cpp index 9d39ec2f9..6475d1463 100644 --- a/Source/Particles/MultiParticleContainer.cpp +++ b/Source/Particles/MultiParticleContainer.cpp @@ -297,7 +297,7 @@ MultiParticleContainer::GetChargeDensity (int lev, bool local) std::unique_ptr<MultiFab> rho = allcontainers[0]->GetChargeDensity(lev, true); for (unsigned i = 1, n = allcontainers.size(); i < n; ++i) { std::unique_ptr<MultiFab> rhoi = allcontainers[i]->GetChargeDensity(lev, true); - MultiFab::Add(*rho, *rhoi, 0, 0, 1, rho->nGrow()); + MultiFab::Add(*rho, *rhoi, 0, 0, rho->nComp(), rho->nGrow()); } if (!local) { const Geometry& gm = allcontainers[0]->Geom(lev); diff --git a/Source/Particles/PhysicalParticleContainer.cpp b/Source/Particles/PhysicalParticleContainer.cpp index 43b46ec49..93a0ad7ea 100644 --- a/Source/Particles/PhysicalParticleContainer.cpp +++ b/Source/Particles/PhysicalParticleContainer.cpp @@ -41,16 +41,19 @@ NumParticlesToAdd(const Box& overlap_box, const RealBox& overlap_realbox, int ref_num_ppc = num_ppc * AMREX_D_TERM(fac, *fac, *fac); for (int i_part=0; i_part<ref_num_ppc;i_part++) { - std::array<Real, 3> r; - plasma_injector->getPositionUnitBox(r, i_part, fac); + std::array<Real, 3> point; + plasma_injector->getPositionUnitBox(point, i_part, fac); + Real x = overlap_corner[0] + (iv[0] + point[0])*dx[0]; #if ( AMREX_SPACEDIM == 3 ) - Real x = overlap_corner[0] + (iv[0] + r[0])*dx[0]; - Real y = overlap_corner[1] + (iv[1] + r[1])*dx[1]; - Real z = overlap_corner[2] + (iv[2] + r[2])*dx[2]; + Real y = overlap_corner[1] + (iv[1] + point[1])*dx[1]; + Real z = overlap_corner[2] + (iv[2] + point[2])*dx[2]; #elif ( AMREX_SPACEDIM == 2 ) - Real x = overlap_corner[0] + (iv[0] + r[0])*dx[0]; Real y = 0; - Real z = overlap_corner[1] + (iv[1] + r[1])*dx[1]; +#ifdef WARPX_RZ + Real z = overlap_corner[1] + (iv[1] + point[2])*dx[1]; +#else + Real z = overlap_corner[1] + (iv[1] + point[1])*dx[1]; +#endif #endif // If the new particle is not inside the tile box, // go to the next generated particle. @@ -448,16 +451,19 @@ PhysicalParticleContainer::AddPlasmaCPU (int lev, RealBox part_realbox) int ref_num_ppc = num_ppc * AMREX_D_TERM(fac, *fac, *fac); for (int i_part=0; i_part<ref_num_ppc;i_part++) { - std::array<Real, 3> r; - plasma_injector->getPositionUnitBox(r, i_part, fac); + std::array<Real, 3> point; + plasma_injector->getPositionUnitBox(point, i_part, fac); + Real x = overlap_corner[0] + (iv[0] + point[0])*dx[0]; #if ( AMREX_SPACEDIM == 3 ) - Real x = overlap_corner[0] + (iv[0] + r[0])*dx[0]; - Real y = overlap_corner[1] + (iv[1] + r[1])*dx[1]; - Real z = overlap_corner[2] + (iv[2] + r[2])*dx[2]; + Real y = overlap_corner[1] + (iv[1] + point[1])*dx[1]; + Real z = overlap_corner[2] + (iv[2] + point[2])*dx[2]; #elif ( AMREX_SPACEDIM == 2 ) - Real x = overlap_corner[0] + (iv[0] + r[0])*dx[0]; Real y = 0; - Real z = overlap_corner[1] + (iv[1] + r[1])*dx[1]; +#ifdef WARPX_RZ + Real z = overlap_corner[1] + (iv[1] + point[2])*dx[1]; +#else + Real z = overlap_corner[1] + (iv[1] + point[1])*dx[1]; +#endif #endif // If the new particle is not inside the tile box, // go to the next generated particle. @@ -473,11 +479,18 @@ PhysicalParticleContainer::AddPlasmaCPU (int lev, RealBox part_realbox) Real yb = y; #ifdef WARPX_RZ - // Replace the x and y, choosing the angle randomly. + // Replace the x and y, setting an angle theta. // These x and y are used to get the momentum and density - Real theta = 2.*MathConst::pi*amrex::Random(); - y = x*std::sin(theta); - x = x*std::cos(theta); + Real theta; + if (WarpX::nmodes == 1) { + // With only 1 mode, the angle doesn't matter so + // choose it randomly. + theta = 2.*MathConst::pi*amrex::Random(); + } else { + theta = 2.*MathConst::pi*point[1]; + } + x = xb*std::cos(theta); + y = xb*std::sin(theta); #endif Real dens; @@ -690,16 +703,19 @@ PhysicalParticleContainer::AddPlasmaGPU (int lev, RealBox part_realbox) int ref_num_ppc = num_ppc * AMREX_D_TERM(fac, *fac, *fac); for (int i_part=0; i_part<ref_num_ppc;i_part++) { - std::array<Real, 3> r; - plasma_injector->getPositionUnitBox(r, i_part, fac); + std::array<Real, 3> point; + plasma_injector->getPositionUnitBox(point, i_part, fac); + Real x = overlap_corner[0] + (iv[0] + point[0])*dx[0]; #if ( AMREX_SPACEDIM == 3 ) - Real x = overlap_corner[0] + (iv[0] + r[0])*dx[0]; - Real y = overlap_corner[1] + (iv[1] + r[1])*dx[1]; - Real z = overlap_corner[2] + (iv[2] + r[2])*dx[2]; + Real y = overlap_corner[1] + (iv[1] + point[1])*dx[1]; + Real z = overlap_corner[2] + (iv[2] + point[2])*dx[2]; #elif ( AMREX_SPACEDIM == 2 ) - Real x = overlap_corner[0] + (iv[0] + r[0])*dx[0]; Real y = 0; - Real z = overlap_corner[1] + (iv[1] + r[1])*dx[1]; +#ifdef WARPX_RZ + Real z = overlap_corner[1] + (iv[1] + point[2])*dx[1]; +#else + Real z = overlap_corner[1] + (iv[1] + point[1])*dx[1]; +#endif #endif // If the new particle is not inside the tile box, // go to the next generated particle. @@ -715,9 +731,16 @@ PhysicalParticleContainer::AddPlasmaGPU (int lev, RealBox part_realbox) Real yb = y; #ifdef WARPX_RZ - // Replace the x and y, choosing the angle randomly. + // Replace the x and y, setting an angle theta. // These x and y are used to get the momentum and density - Real theta = 2.*MathConst::pi*amrex::Random(); + Real theta; + if (WarpX::nmodes == 1) { + // With only 1 mode, the angle doesn't matter so + // choose it randomly. + theta = 2.*MathConst::pi*amrex::Random(); + } else { + theta = 2.*MathConst::pi*point[1]; + } x = xb*std::cos(theta); y = xb*std::sin(theta); #endif @@ -1133,6 +1156,7 @@ PhysicalParticleContainer::FieldGather (int lev, BL_TO_FORTRAN_ANYD(bxfab), BL_TO_FORTRAN_ANYD(byfab), BL_TO_FORTRAN_ANYD(bzfab), + &WarpX::nmodes, &ll4symtry, &WarpX::l_lower_order_in_v, &WarpX::do_nodal, &lvect_fieldgathe, &WarpX::field_gathering_algo); @@ -1425,6 +1449,7 @@ PhysicalParticleContainer::Evolve (int lev, BL_TO_FORTRAN_ANYD(*bxfab), BL_TO_FORTRAN_ANYD(*byfab), BL_TO_FORTRAN_ANYD(*bzfab), + &WarpX::nmodes, &ll4symtry, &WarpX::l_lower_order_in_v, &WarpX::do_nodal, &lvect_fieldgathe, &WarpX::field_gathering_algo); @@ -1512,6 +1537,7 @@ PhysicalParticleContainer::Evolve (int lev, BL_TO_FORTRAN_ANYD(*cbxfab), BL_TO_FORTRAN_ANYD(*cbyfab), BL_TO_FORTRAN_ANYD(*cbzfab), + &WarpX::nmodes, &ll4symtry, &WarpX::l_lower_order_in_v, &WarpX::do_nodal, &lvect_fieldgathe, &WarpX::field_gathering_algo); } @@ -1853,6 +1879,7 @@ PhysicalParticleContainer::PushP (int lev, Real dt, BL_TO_FORTRAN_ANYD(bxfab), BL_TO_FORTRAN_ANYD(byfab), BL_TO_FORTRAN_ANYD(bzfab), + &WarpX::nmodes, &ll4symtry, &WarpX::l_lower_order_in_v, &WarpX::do_nodal, &lvect_fieldgathe, &WarpX::field_gathering_algo); diff --git a/Source/Particles/RigidInjectedParticleContainer.cpp b/Source/Particles/RigidInjectedParticleContainer.cpp index 2a3e8dd0d..d659b3854 100644 --- a/Source/Particles/RigidInjectedParticleContainer.cpp +++ b/Source/Particles/RigidInjectedParticleContainer.cpp @@ -426,6 +426,7 @@ RigidInjectedParticleContainer::PushP (int lev, Real dt, BL_TO_FORTRAN_ANYD(bxfab), BL_TO_FORTRAN_ANYD(byfab), BL_TO_FORTRAN_ANYD(bzfab), + &WarpX::nmodes, &ll4symtry, &l_lower_order_in_v, &WarpX::do_nodal, &lvect_fieldgathe, &WarpX::field_gathering_algo); diff --git a/Source/Particles/WarpXParticleContainer.H b/Source/Particles/WarpXParticleContainer.H index 0e800bf1d..b2d86c587 100644 --- a/Source/Particles/WarpXParticleContainer.H +++ b/Source/Particles/WarpXParticleContainer.H @@ -42,6 +42,9 @@ namespace ParticleStringNames {"Bx", PIdx::Bx }, {"By", PIdx::By }, {"Bz", PIdx::Bz } +#ifdef WARPX_RZ + ,{"theta", PIdx::theta} +#endif }; } diff --git a/Source/Particles/WarpXParticleContainer.cpp b/Source/Particles/WarpXParticleContainer.cpp index 317f46fd4..a1517cb44 100644 --- a/Source/Particles/WarpXParticleContainer.cpp +++ b/Source/Particles/WarpXParticleContainer.cpp @@ -354,9 +354,9 @@ WarpXParticleContainer::DepositCurrent(WarpXParIter& pti, tby.grow(ngJ); tbz.grow(ngJ); - local_jx[thread_num].resize(tbx); - local_jy[thread_num].resize(tby); - local_jz[thread_num].resize(tbz); + local_jx[thread_num].resize(tbx, jx->nComp()); + local_jy[thread_num].resize(tby, jy->nComp()); + local_jz[thread_num].resize(tbz, jz->nComp()); Real* jx_ptr = local_jx[thread_num].dataPtr(); Real* jy_ptr = local_jy[thread_num].dataPtr(); @@ -379,6 +379,7 @@ WarpXParticleContainer::DepositCurrent(WarpXParIter& pti, jx_ptr, &ngJ, jxntot.getVect(), jy_ptr, &ngJ, jyntot.getVect(), jz_ptr, &ngJ, jzntot.getVect(), + &WarpX::nmodes, &np_to_depose, m_xp[thread_num].dataPtr() + offset, m_yp[thread_num].dataPtr() + offset, @@ -399,6 +400,7 @@ WarpXParticleContainer::DepositCurrent(WarpXParIter& pti, jx_ptr, &ngJ, jxntot.getVect(), jy_ptr, &ngJ, jyntot.getVect(), jz_ptr, &ngJ, jzntot.getVect(), + &WarpX::nmodes, &xyzmin[0], &dx[0]); #endif BL_PROFILE_VAR_STOP(blp_pxr_cd); @@ -407,9 +409,9 @@ WarpXParticleContainer::DepositCurrent(WarpXParIter& pti, BL_PROFILE_VAR_START(blp_accumulate); // CPU, tiling: atomicAdd local_jx into jx // (same for jx and jz) - (*jx)[pti].atomicAdd(local_jx[thread_num], tbx, tbx, 0, 0, 1); - (*jy)[pti].atomicAdd(local_jy[thread_num], tby, tby, 0, 0, 1); - (*jz)[pti].atomicAdd(local_jz[thread_num], tbz, tbz, 0, 0, 1); + (*jx)[pti].atomicAdd(local_jx[thread_num], tbx, tbx, 0, 0, local_jx[thread_num].nComp()); + (*jy)[pti].atomicAdd(local_jy[thread_num], tby, tby, 0, 0, local_jy[thread_num].nComp()); + (*jz)[pti].atomicAdd(local_jz[thread_num], tbz, tbz, 0, 0, local_jz[thread_num].nComp()); BL_PROFILE_VAR_STOP(blp_accumulate); #endif } @@ -440,11 +442,11 @@ WarpXParticleContainer::DepositCharge ( WarpXParIter& pti, RealVector& wp, const std::array<Real, 3>& xyzmin = xyzmin_tile; #ifdef AMREX_USE_GPU - data_ptr = (*rhomf)[pti].dataPtr(icomp); + data_ptr = (*rhomf)[pti].dataPtr(icomp*(rhomf->nComp()/2)); auto rholen = (*rhomf)[pti].length(); #else tile_box.grow(ngRho); - local_rho[thread_num].resize(tile_box); + local_rho[thread_num].resize(tile_box, rhomf->nComp()); data_ptr = local_rho[thread_num].dataPtr(); auto rholen = local_rho[thread_num].length(); @@ -483,7 +485,7 @@ WarpXParticleContainer::DepositCharge ( WarpXParIter& pti, RealVector& wp, #ifndef AMREX_USE_GPU BL_PROFILE_VAR_START(blp_accumulate); - (*rhomf)[pti].atomicAdd(local_rho[thread_num], tile_box, tile_box, 0, icomp, 1); + (*rhomf)[pti].atomicAdd(local_rho[thread_num], tile_box, tile_box, 0, icomp*(rhomf->nComp()/2), (rhomf->nComp()/2)); BL_PROFILE_VAR_STOP(blp_accumulate); #endif @@ -502,7 +504,7 @@ WarpXParticleContainer::DepositCharge ( WarpXParIter& pti, RealVector& wp, #else tile_box = amrex::convert(ctilebox, IntVect::TheUnitVector()); tile_box.grow(ngRho); - local_rho[thread_num].resize(tile_box); + local_rho[thread_num].resize(tile_box, crhomf->nComp()); data_ptr = local_rho[thread_num].dataPtr(); auto rholen = local_rho[thread_num].length(); @@ -543,7 +545,7 @@ WarpXParticleContainer::DepositCharge ( WarpXParIter& pti, RealVector& wp, #ifndef AMREX_USE_GPU BL_PROFILE_VAR_START(blp_accumulate); - (*crhomf)[pti].atomicAdd(local_rho[thread_num], tile_box, tile_box, 0, icomp, 1); + (*crhomf)[pti].atomicAdd(local_rho[thread_num], tile_box, tile_box, 0, icomp*(crhomf->nComp()/2), (crhomf->nComp()/2)); BL_PROFILE_VAR_STOP(blp_accumulate); #endif @@ -598,7 +600,7 @@ WarpXParticleContainer::DepositCharge (Vector<std::unique_ptr<MultiFab> >& rho, BoxArray coarsened_fine_BA = fine_BA; coarsened_fine_BA.coarsen(m_gdb->refRatio(lev)); - MultiFab coarsened_fine_data(coarsened_fine_BA, fine_dm, 1, 0); + MultiFab coarsened_fine_data(coarsened_fine_BA, fine_dm, rho[lev+1]->nComp(), 0); coarsened_fine_data.setVal(0.0); IntVect ratio(AMREX_D_DECL(2, 2, 2)); // FIXME @@ -629,7 +631,7 @@ WarpXParticleContainer::GetChargeDensity (int lev, bool local) const int ng = WarpX::nox; - auto rho = std::unique_ptr<MultiFab>(new MultiFab(nba,dm,1,ng)); + auto rho = std::unique_ptr<MultiFab>(new MultiFab(nba,dm,WarpX::ncomps,ng)); rho->setVal(0.0); #ifdef _OPENMP @@ -657,7 +659,7 @@ WarpXParticleContainer::GetChargeDensity (int lev, bool local) Box tile_box = convert(pti.tilebox(), IntVect::TheUnitVector()); const std::array<Real, 3>& xyzmin = xyzmin_tile; tile_box.grow(ng); - rho_loc.resize(tile_box); + rho_loc.resize(tile_box, rho->nComp()); rho_loc = 0.0; data_ptr = rho_loc.dataPtr(); auto rholen = rho_loc.length(); diff --git a/Source/Python/WarpXWrappers.cpp b/Source/Python/WarpXWrappers.cpp index 3c1a930b3..3ed4830f5 100644 --- a/Source/Python/WarpXWrappers.cpp +++ b/Source/Python/WarpXWrappers.cpp @@ -10,22 +10,26 @@ namespace { - double** getMultiFabPointers(const amrex::MultiFab& mf, int *num_boxes, int *ngrow, int **shapes) + double** getMultiFabPointers(const amrex::MultiFab& mf, int *num_boxes, int *ncomps, int *ngrow, int **shapes) { + *ncomps = mf.nComp(); *ngrow = mf.nGrow(); *num_boxes = mf.local_size(); - *shapes = (int*) malloc(AMREX_SPACEDIM * (*num_boxes) * sizeof(int)); + int shapesize = AMREX_SPACEDIM; + if (mf.nComp() > 1) shapesize += 1; + *shapes = (int*) malloc(shapesize * (*num_boxes) * sizeof(int)); double** data = (double**) malloc((*num_boxes) * sizeof(double*)); - int i = 0; #ifdef _OPENMP #pragma omp parallel #endif - for ( amrex::MFIter mfi(mf, false); mfi.isValid(); ++mfi, ++i ) { + for ( amrex::MFIter mfi(mf, false); mfi.isValid(); ++mfi ) { + int i = mfi.LocalIndex(); data[i] = (double*) mf[mfi].dataPtr(); for (int j = 0; j < AMREX_SPACEDIM; ++j) { - (*shapes)[AMREX_SPACEDIM*i+j] = mf[mfi].box().length(j); + (*shapes)[shapesize*i+j] = mf[mfi].box().length(j); } + if (mf.nComp() > 1) (*shapes)[shapesize*i+2] = mf.nComp(); } return data; } @@ -197,9 +201,9 @@ extern "C" } double** warpx_getEfield(int lev, int direction, - int *return_size, int *ngrow, int **shapes) { + int *return_size, int *ncomps, int *ngrow, int **shapes) { auto & mf = WarpX::GetInstance().getEfield(lev, direction); - return getMultiFabPointers(mf, return_size, ngrow, shapes); + return getMultiFabPointers(mf, return_size, ncomps, ngrow, shapes); } int* warpx_getEfieldLoVects(int lev, int direction, @@ -209,9 +213,9 @@ extern "C" } double** warpx_getEfieldCP(int lev, int direction, - int *return_size, int *ngrow, int **shapes) { + int *return_size, int *ncomps, int *ngrow, int **shapes) { auto & mf = WarpX::GetInstance().getEfield_cp(lev, direction); - return getMultiFabPointers(mf, return_size, ngrow, shapes); + return getMultiFabPointers(mf, return_size, ncomps, ngrow, shapes); } int* warpx_getEfieldCPLoVects(int lev, int direction, @@ -221,9 +225,9 @@ extern "C" } double** warpx_getEfieldFP(int lev, int direction, - int *return_size, int *ngrow, int **shapes) { + int *return_size, int *ncomps, int *ngrow, int **shapes) { auto & mf = WarpX::GetInstance().getEfield_fp(lev, direction); - return getMultiFabPointers(mf, return_size, ngrow, shapes); + return getMultiFabPointers(mf, return_size, ncomps, ngrow, shapes); } int* warpx_getEfieldFPLoVects(int lev, int direction, @@ -233,9 +237,9 @@ extern "C" } double** warpx_getBfield(int lev, int direction, - int *return_size, int *ngrow, int **shapes) { + int *return_size, int *ncomps, int *ngrow, int **shapes) { auto & mf = WarpX::GetInstance().getBfield(lev, direction); - return getMultiFabPointers(mf, return_size, ngrow, shapes); + return getMultiFabPointers(mf, return_size, ncomps, ngrow, shapes); } int* warpx_getBfieldLoVects(int lev, int direction, @@ -245,9 +249,9 @@ extern "C" } double** warpx_getBfieldCP(int lev, int direction, - int *return_size, int *ngrow, int **shapes) { + int *return_size, int *ncomps, int *ngrow, int **shapes) { auto & mf = WarpX::GetInstance().getBfield_cp(lev, direction); - return getMultiFabPointers(mf, return_size, ngrow, shapes); + return getMultiFabPointers(mf, return_size, ncomps, ngrow, shapes); } int* warpx_getBfieldCPLoVects(int lev, int direction, @@ -257,9 +261,9 @@ extern "C" } double** warpx_getBfieldFP(int lev, int direction, - int *return_size, int *ngrow, int **shapes) { + int *return_size, int *ncomps, int *ngrow, int **shapes) { auto & mf = WarpX::GetInstance().getBfield_fp(lev, direction); - return getMultiFabPointers(mf, return_size, ngrow, shapes); + return getMultiFabPointers(mf, return_size, ncomps, ngrow, shapes); } int* warpx_getBfieldFPLoVects(int lev, int direction, @@ -269,9 +273,9 @@ extern "C" } double** warpx_getCurrentDensity(int lev, int direction, - int *return_size, int *ngrow, int **shapes) { + int *return_size, int *ncomps, int *ngrow, int **shapes) { auto & mf = WarpX::GetInstance().getcurrent(lev, direction); - return getMultiFabPointers(mf, return_size, ngrow, shapes); + return getMultiFabPointers(mf, return_size, ncomps, ngrow, shapes); } int* warpx_getCurrentDensityLoVects(int lev, int direction, @@ -281,9 +285,9 @@ extern "C" } double** warpx_getCurrentDensityCP(int lev, int direction, - int *return_size, int *ngrow, int **shapes) { + int *return_size, int *ncomps, int *ngrow, int **shapes) { auto & mf = WarpX::GetInstance().getcurrent_cp(lev, direction); - return getMultiFabPointers(mf, return_size, ngrow, shapes); + return getMultiFabPointers(mf, return_size, ncomps, ngrow, shapes); } int* warpx_getCurrentDensityCPLoVects(int lev, int direction, @@ -293,9 +297,9 @@ extern "C" } double** warpx_getCurrentDensityFP(int lev, int direction, - int *return_size, int *ngrow, int **shapes) { + int *return_size, int *ncomps, int *ngrow, int **shapes) { auto & mf = WarpX::GetInstance().getcurrent_fp(lev, direction); - return getMultiFabPointers(mf, return_size, ngrow, shapes); + return getMultiFabPointers(mf, return_size, ncomps, ngrow, shapes); } int* warpx_getCurrentDensityFPLoVects(int lev, int direction, diff --git a/Source/Python/WarpXWrappers.h b/Source/Python/WarpXWrappers.h index 94fbb0d30..44e0ed4e1 100644 --- a/Source/Python/WarpXWrappers.h +++ b/Source/Python/WarpXWrappers.h @@ -62,19 +62,19 @@ extern "C" { long warpx_getNumParticles(int speciesnumber); double** warpx_getEfield(int lev, int direction, - int *return_size, int* ngrow, int **shapes); + int *return_size, int* ncomps, int* ngrow, int **shapes); int* warpx_getEfieldLoVects(int lev, int direction, int *return_size, int* ngrow); double** warpx_getBfield(int lev, int direction, - int *return_size, int* ngrow, int **shapes); + int *return_size, int* ncomps, int* ngrow, int **shapes); int* warpx_getBfieldLoVects(int lev, int direction, int *return_size, int* ngrow); double** warpx_getCurrentDensity(int lev, int direction, - int *return_size, int* ngrow, int **shapes); + int *return_size, int* ncomps, int* ngrow, int **shapes); int* warpx_getCurrentDensityLoVects(int lev, int direction, int *return_size, int* ngrow); diff --git a/Source/WarpX.H b/Source/WarpX.H index 4ad3d119f..5d3008972 100644 --- a/Source/WarpX.H +++ b/Source/WarpX.H @@ -92,6 +92,10 @@ public: static long noy; static long noz; + // Number of modes for the RZ multimode version + static long nmodes; + static long ncomps; + static bool use_fdtd_nci_corr; static int l_lower_order_in_v; diff --git a/Source/WarpX.cpp b/Source/WarpX.cpp index 877882037..4ad2cf75d 100644 --- a/Source/WarpX.cpp +++ b/Source/WarpX.cpp @@ -43,6 +43,9 @@ long WarpX::field_gathering_algo; long WarpX::particle_pusher_algo; int WarpX::maxwell_fdtd_solver_id; +long WarpX::nmodes = 1; +long WarpX::ncomps = 1; + long WarpX::nox = 1; long WarpX::noy = 1; long WarpX::noz = 1; @@ -470,6 +473,10 @@ WarpX::ReadParameters () // Use same shape factors in all directions, for gathering l_lower_order_in_v = false; } + + // Only needs to be set with WARPX_RZ, otherwise defaults to 1. + pp.query("nmodes", nmodes); + } { @@ -700,20 +707,31 @@ void WarpX::AllocLevelMFs (int lev, const BoxArray& ba, const DistributionMapping& dm, const IntVect& ngE, const IntVect& ngJ, const IntVect& ngRho, int ngF) { + +#if defined WARPX_RZ + if (nmodes > 1) { + // There is a real and imaginary component for each mode + ncomps = nmodes*2; + } else { + // With only mode 0, only reals are used + ncomps = 1; + } +#endif + // // The fine patch // - Bfield_fp[lev][0].reset( new MultiFab(amrex::convert(ba,Bx_nodal_flag),dm,1,ngE)); - Bfield_fp[lev][1].reset( new MultiFab(amrex::convert(ba,By_nodal_flag),dm,1,ngE)); - Bfield_fp[lev][2].reset( new MultiFab(amrex::convert(ba,Bz_nodal_flag),dm,1,ngE)); + Bfield_fp[lev][0].reset( new MultiFab(amrex::convert(ba,Bx_nodal_flag),dm,ncomps,ngE)); + Bfield_fp[lev][1].reset( new MultiFab(amrex::convert(ba,By_nodal_flag),dm,ncomps,ngE)); + Bfield_fp[lev][2].reset( new MultiFab(amrex::convert(ba,Bz_nodal_flag),dm,ncomps,ngE)); - Efield_fp[lev][0].reset( new MultiFab(amrex::convert(ba,Ex_nodal_flag),dm,1,ngE)); - Efield_fp[lev][1].reset( new MultiFab(amrex::convert(ba,Ey_nodal_flag),dm,1,ngE)); - Efield_fp[lev][2].reset( new MultiFab(amrex::convert(ba,Ez_nodal_flag),dm,1,ngE)); + Efield_fp[lev][0].reset( new MultiFab(amrex::convert(ba,Ex_nodal_flag),dm,ncomps,ngE)); + Efield_fp[lev][1].reset( new MultiFab(amrex::convert(ba,Ey_nodal_flag),dm,ncomps,ngE)); + Efield_fp[lev][2].reset( new MultiFab(amrex::convert(ba,Ez_nodal_flag),dm,ncomps,ngE)); - current_fp[lev][0].reset( new MultiFab(amrex::convert(ba,jx_nodal_flag),dm,1,ngJ)); - current_fp[lev][1].reset( new MultiFab(amrex::convert(ba,jy_nodal_flag),dm,1,ngJ)); - current_fp[lev][2].reset( new MultiFab(amrex::convert(ba,jz_nodal_flag),dm,1,ngJ)); + current_fp[lev][0].reset( new MultiFab(amrex::convert(ba,jx_nodal_flag),dm,ncomps,ngJ)); + current_fp[lev][1].reset( new MultiFab(amrex::convert(ba,jy_nodal_flag),dm,ncomps,ngJ)); + current_fp[lev][2].reset( new MultiFab(amrex::convert(ba,jz_nodal_flag),dm,ncomps,ngJ)); const auto& period = Geom(lev).periodicity(); current_fp_owner_masks[lev][0] = std::move(current_fp[lev][0]->OwnerMask(period)); @@ -722,25 +740,25 @@ WarpX::AllocLevelMFs (int lev, const BoxArray& ba, const DistributionMapping& dm if (do_dive_cleaning || plot_rho) { - rho_fp[lev].reset(new MultiFab(amrex::convert(ba,IntVect::TheUnitVector()),dm,2,ngRho)); + rho_fp[lev].reset(new MultiFab(amrex::convert(ba,IntVect::TheUnitVector()),dm,2*ncomps,ngRho)); rho_fp_owner_masks[lev] = std::move(rho_fp[lev]->OwnerMask(period)); } if (do_subcycling == 1 && lev == 0) { - current_store[lev][0].reset( new MultiFab(amrex::convert(ba,jx_nodal_flag),dm,1,ngJ)); - current_store[lev][1].reset( new MultiFab(amrex::convert(ba,jy_nodal_flag),dm,1,ngJ)); - current_store[lev][2].reset( new MultiFab(amrex::convert(ba,jz_nodal_flag),dm,1,ngJ)); + current_store[lev][0].reset( new MultiFab(amrex::convert(ba,jx_nodal_flag),dm,ncomps,ngJ)); + current_store[lev][1].reset( new MultiFab(amrex::convert(ba,jy_nodal_flag),dm,ncomps,ngJ)); + current_store[lev][2].reset( new MultiFab(amrex::convert(ba,jz_nodal_flag),dm,ncomps,ngJ)); } if (do_dive_cleaning) { - F_fp[lev].reset (new MultiFab(amrex::convert(ba,IntVect::TheUnitVector()),dm,1, ngF)); + F_fp[lev].reset (new MultiFab(amrex::convert(ba,IntVect::TheUnitVector()),dm,ncomps, ngF)); } #ifdef WARPX_USE_PSATD else { - rho_fp[lev].reset(new MultiFab(amrex::convert(ba,IntVect::TheUnitVector()),dm,2,ngRho)); + rho_fp[lev].reset(new MultiFab(amrex::convert(ba,IntVect::TheUnitVector()),dm,2*ncomps,ngRho)); rho_fp_owner_masks[lev] = std::move(rho_fp[lev]->OwnerMask(period)); } #endif @@ -751,19 +769,19 @@ WarpX::AllocLevelMFs (int lev, const BoxArray& ba, const DistributionMapping& dm if (lev == 0) { for (int idir = 0; idir < 3; ++idir) { - Efield_aux[lev][idir].reset(new MultiFab(*Efield_fp[lev][idir], amrex::make_alias, 0, 1)); - Bfield_aux[lev][idir].reset(new MultiFab(*Bfield_fp[lev][idir], amrex::make_alias, 0, 1)); + Efield_aux[lev][idir].reset(new MultiFab(*Efield_fp[lev][idir], amrex::make_alias, 0, ncomps)); + Bfield_aux[lev][idir].reset(new MultiFab(*Bfield_fp[lev][idir], amrex::make_alias, 0, ncomps)); } } else { - Bfield_aux[lev][0].reset( new MultiFab(amrex::convert(ba,Bx_nodal_flag),dm,1,ngE)); - Bfield_aux[lev][1].reset( new MultiFab(amrex::convert(ba,By_nodal_flag),dm,1,ngE)); - Bfield_aux[lev][2].reset( new MultiFab(amrex::convert(ba,Bz_nodal_flag),dm,1,ngE)); + Bfield_aux[lev][0].reset( new MultiFab(amrex::convert(ba,Bx_nodal_flag),dm,ncomps,ngE)); + Bfield_aux[lev][1].reset( new MultiFab(amrex::convert(ba,By_nodal_flag),dm,ncomps,ngE)); + Bfield_aux[lev][2].reset( new MultiFab(amrex::convert(ba,Bz_nodal_flag),dm,ncomps,ngE)); - Efield_aux[lev][0].reset( new MultiFab(amrex::convert(ba,Ex_nodal_flag),dm,1,ngE)); - Efield_aux[lev][1].reset( new MultiFab(amrex::convert(ba,Ey_nodal_flag),dm,1,ngE)); - Efield_aux[lev][2].reset( new MultiFab(amrex::convert(ba,Ez_nodal_flag),dm,1,ngE)); + Efield_aux[lev][0].reset( new MultiFab(amrex::convert(ba,Ex_nodal_flag),dm,ncomps,ngE)); + Efield_aux[lev][1].reset( new MultiFab(amrex::convert(ba,Ey_nodal_flag),dm,ncomps,ngE)); + Efield_aux[lev][2].reset( new MultiFab(amrex::convert(ba,Ez_nodal_flag),dm,ncomps,ngE)); } // @@ -775,19 +793,19 @@ WarpX::AllocLevelMFs (int lev, const BoxArray& ba, const DistributionMapping& dm cba.coarsen(refRatio(lev-1)); // Create the MultiFabs for B - Bfield_cp[lev][0].reset( new MultiFab(amrex::convert(cba,Bx_nodal_flag),dm,1,ngE)); - Bfield_cp[lev][1].reset( new MultiFab(amrex::convert(cba,By_nodal_flag),dm,1,ngE)); - Bfield_cp[lev][2].reset( new MultiFab(amrex::convert(cba,Bz_nodal_flag),dm,1,ngE)); + Bfield_cp[lev][0].reset( new MultiFab(amrex::convert(cba,Bx_nodal_flag),dm,ncomps,ngE)); + Bfield_cp[lev][1].reset( new MultiFab(amrex::convert(cba,By_nodal_flag),dm,ncomps,ngE)); + Bfield_cp[lev][2].reset( new MultiFab(amrex::convert(cba,Bz_nodal_flag),dm,ncomps,ngE)); // Create the MultiFabs for E - Efield_cp[lev][0].reset( new MultiFab(amrex::convert(cba,Ex_nodal_flag),dm,1,ngE)); - Efield_cp[lev][1].reset( new MultiFab(amrex::convert(cba,Ey_nodal_flag),dm,1,ngE)); - Efield_cp[lev][2].reset( new MultiFab(amrex::convert(cba,Ez_nodal_flag),dm,1,ngE)); + Efield_cp[lev][0].reset( new MultiFab(amrex::convert(cba,Ex_nodal_flag),dm,ncomps,ngE)); + Efield_cp[lev][1].reset( new MultiFab(amrex::convert(cba,Ey_nodal_flag),dm,ncomps,ngE)); + Efield_cp[lev][2].reset( new MultiFab(amrex::convert(cba,Ez_nodal_flag),dm,ncomps,ngE)); // Create the MultiFabs for the current - current_cp[lev][0].reset( new MultiFab(amrex::convert(cba,jx_nodal_flag),dm,1,ngJ)); - current_cp[lev][1].reset( new MultiFab(amrex::convert(cba,jy_nodal_flag),dm,1,ngJ)); - current_cp[lev][2].reset( new MultiFab(amrex::convert(cba,jz_nodal_flag),dm,1,ngJ)); + current_cp[lev][0].reset( new MultiFab(amrex::convert(cba,jx_nodal_flag),dm,ncomps,ngJ)); + current_cp[lev][1].reset( new MultiFab(amrex::convert(cba,jy_nodal_flag),dm,ncomps,ngJ)); + current_cp[lev][2].reset( new MultiFab(amrex::convert(cba,jz_nodal_flag),dm,ncomps,ngJ)); const auto& cperiod = Geom(lev-1).periodicity(); current_cp_owner_masks[lev][0] = std::move(current_cp[lev][0]->OwnerMask(cperiod)); @@ -795,17 +813,17 @@ WarpX::AllocLevelMFs (int lev, const BoxArray& ba, const DistributionMapping& dm current_cp_owner_masks[lev][2] = std::move(current_cp[lev][2]->OwnerMask(cperiod)); if (do_dive_cleaning || plot_rho){ - rho_cp[lev].reset(new MultiFab(amrex::convert(cba,IntVect::TheUnitVector()),dm,2,ngRho)); + rho_cp[lev].reset(new MultiFab(amrex::convert(cba,IntVect::TheUnitVector()),dm,2*ncomps,ngRho)); rho_cp_owner_masks[lev] = std::move(rho_cp[lev]->OwnerMask(cperiod)); } if (do_dive_cleaning) { - F_cp[lev].reset (new MultiFab(amrex::convert(cba,IntVect::TheUnitVector()),dm,1, ngF)); + F_cp[lev].reset (new MultiFab(amrex::convert(cba,IntVect::TheUnitVector()),dm,ncomps, ngF)); } #ifdef WARPX_USE_PSATD else { - rho_cp[lev].reset(new MultiFab(amrex::convert(cba,IntVect::TheUnitVector()),dm,2,ngRho)); + rho_cp[lev].reset(new MultiFab(amrex::convert(cba,IntVect::TheUnitVector()),dm,2*ncomps,ngRho)); rho_cp_owner_masks[lev] = std::move(rho_cp[lev]->OwnerMask(cperiod)); } #endif @@ -821,28 +839,28 @@ WarpX::AllocLevelMFs (int lev, const BoxArray& ba, const DistributionMapping& dm if (n_field_gather_buffer > 0) { // Create the MultiFabs for B - Bfield_cax[lev][0].reset( new MultiFab(amrex::convert(cba,Bx_nodal_flag),dm,1,ngE)); - Bfield_cax[lev][1].reset( new MultiFab(amrex::convert(cba,By_nodal_flag),dm,1,ngE)); - Bfield_cax[lev][2].reset( new MultiFab(amrex::convert(cba,Bz_nodal_flag),dm,1,ngE)); + Bfield_cax[lev][0].reset( new MultiFab(amrex::convert(cba,Bx_nodal_flag),dm,ncomps,ngE)); + Bfield_cax[lev][1].reset( new MultiFab(amrex::convert(cba,By_nodal_flag),dm,ncomps,ngE)); + Bfield_cax[lev][2].reset( new MultiFab(amrex::convert(cba,Bz_nodal_flag),dm,ncomps,ngE)); // Create the MultiFabs for E - Efield_cax[lev][0].reset( new MultiFab(amrex::convert(cba,Ex_nodal_flag),dm,1,ngE)); - Efield_cax[lev][1].reset( new MultiFab(amrex::convert(cba,Ey_nodal_flag),dm,1,ngE)); - Efield_cax[lev][2].reset( new MultiFab(amrex::convert(cba,Ez_nodal_flag),dm,1,ngE)); + Efield_cax[lev][0].reset( new MultiFab(amrex::convert(cba,Ex_nodal_flag),dm,ncomps,ngE)); + Efield_cax[lev][1].reset( new MultiFab(amrex::convert(cba,Ey_nodal_flag),dm,ncomps,ngE)); + Efield_cax[lev][2].reset( new MultiFab(amrex::convert(cba,Ez_nodal_flag),dm,ncomps,ngE)); - gather_buffer_masks[lev].reset( new iMultiFab(ba, dm, 1, 1) ); + gather_buffer_masks[lev].reset( new iMultiFab(ba, dm, ncomps, 1) ); // Gather buffer masks have 1 ghost cell, because of the fact // that particles may move by more than one cell when using subcycling. } if (n_current_deposition_buffer > 0) { - current_buf[lev][0].reset( new MultiFab(amrex::convert(cba,jx_nodal_flag),dm,1,ngJ)); - current_buf[lev][1].reset( new MultiFab(amrex::convert(cba,jy_nodal_flag),dm,1,ngJ)); - current_buf[lev][2].reset( new MultiFab(amrex::convert(cba,jz_nodal_flag),dm,1,ngJ)); + current_buf[lev][0].reset( new MultiFab(amrex::convert(cba,jx_nodal_flag),dm,ncomps,ngJ)); + current_buf[lev][1].reset( new MultiFab(amrex::convert(cba,jy_nodal_flag),dm,ncomps,ngJ)); + current_buf[lev][2].reset( new MultiFab(amrex::convert(cba,jz_nodal_flag),dm,ncomps,ngJ)); if (do_dive_cleaning || plot_rho) { - charge_buf[lev].reset( new MultiFab(amrex::convert(cba,IntVect::TheUnitVector()),dm,2,ngRho)); + charge_buf[lev].reset( new MultiFab(amrex::convert(cba,IntVect::TheUnitVector()),dm,2*ncomps,ngRho)); } - current_buffer_masks[lev].reset( new iMultiFab(ba, dm, 1, 1) ); + current_buffer_masks[lev].reset( new iMultiFab(ba, dm, ncomps, 1) ); // Current buffer masks have 1 ghost cell, because of the fact // that particles may move by more than one cell when using subcycling. } |