Skip to content

Commit 7348362

Browse files
cemitch99cemitch99
committed
Clean separation of True_2p5D push section.
1 parent 794c06f commit 7348362

File tree

1 file changed

+87
-51
lines changed

1 file changed

+87
-51
lines changed

‎src/particles/spacecharge/GatherAndPush.cpp‎

Lines changed: 87 additions & 51 deletions
Original file line numberDiff line numberDiff line change
@@ -84,50 +84,11 @@ namespace impactx::particles::spacecharge
8484
amrex::ParticleReal const push_consts = dt * charge * inv_gamma2 / pz_ref_SI;
8585

8686
// gather to each particle and push momentum
87-
if (space_charge == SpaceChargeAlgo::True_2D || space_charge == SpaceChargeAlgo::True_2p5D) {
87+
if (space_charge == SpaceChargeAlgo::True_2D) {
8888
// flatten 3rd dimension
8989
auto prob_lo_2D = gm.ProbLoArray();
9090
prob_lo_2D[2] = 0.0_rt;
9191

92-
if (space_charge == SpaceChargeAlgo::True_2p5D) {
93-
// Calculate z-dependent scaling by current
94-
int tp5d_bins = 129;
95-
amrex::ParmParse pp_algo("algo.space_charge");
96-
pp_algo.queryAddWithParser("gauss_charge_z_bins", tp5d_bins);
97-
98-
// Set parameters for charge deposition
99-
bool const is_unity_particle_weight = false;
100-
bool const GetNumberDensity = true;
101-
102-
// Measure beam size, extract the min, max of particle positions
103-
[[maybe_unused]] auto const [x_min, y_min, t_min, x_max, y_max, t_max] =
104-
pc.MinAndMaxPositions();
105-
106-
int const num_bins = tp5d_bins; // Set resolution
107-
amrex::Real const bin_min = t_min;
108-
amrex::Real const bin_max = t_max;
109-
amrex::Real const bin_size = (bin_max - bin_min) / (num_bins - 1); // number of evaluation points
110-
// Allocate memory for the charge profile
111-
amrex::Gpu::DeviceVector<amrex::Real> charge_distribution(num_bins + 1, 0.0);
112-
// Call charge deposition function
113-
impactx::particles::wakefields::DepositCharge1D(pc, charge_distribution, bin_min, bin_size, is_unity_particle_weight);
114-
115-
// Sum up all partial charge histograms to each MPI process to calculate
116-
// the global charge slope.
117-
amrex::ParallelAllReduce::Sum(
118-
charge_distribution.data(),
119-
charge_distribution.size(),
120-
amrex::ParallelDescriptor::Communicator()
121-
);
122-
123-
// Call charge density derivative function
124-
amrex::Gpu::DeviceVector<amrex::Real> slopes(charge_distribution.size() - 1, 0.0);
125-
impactx::particles::wakefields::DerivativeCharge1D(charge_distribution, slopes, bin_size,GetNumberDensity);
126-
127-
[[maybe_unused]] amrex::Real const * const beam_profile_slope = slopes.data();
128-
[[maybe_unused]] amrex::Real const * const beam_profile = charge_distribution.data();
129-
}
130-
13192
amrex::ParallelFor(np, [=] AMREX_GPU_DEVICE (int i) {
13293
// access SoA Real data
13394
amrex::ParticleReal & AMREX_RESTRICT x = part_x[i];
@@ -147,21 +108,96 @@ namespace impactx::particles::spacecharge
147108
);
148109

149110
// push momentum
150-
if (space_charge == SpaceChargeAlgo::True_2D) {
151-
px += field_interp[0] * push_consts * dr[2] / (beta * c0_SI);
152-
py += field_interp[1] * push_consts * dr[2] / (beta * c0_SI);
153-
pz += 0.0_rt;
154-
} else if (space_charge == SpaceChargeAlgo::True_2p5D) {
155-
// TODO: apply z-dependent scaling by current and longitudinal kick
156-
px += field_interp[0] * push_consts;
157-
py += field_interp[1] * push_consts;
158-
pz += 0.0_rt;
159-
} else {
160-
}
111+
px += field_interp[0] * push_consts * dr[2] / (beta * c0_SI);
112+
py += field_interp[1] * push_consts * dr[2] / (beta * c0_SI);
113+
pz += 0.0_rt;
161114

162115
// push position is done in the lattice elements
163116
});
164117
}
118+
if (space_charge == SpaceChargeAlgo::True_2p5D) {
119+
// flatten 3rd dimension
120+
auto prob_lo_2D = gm.ProbLoArray();
121+
prob_lo_2D[2] = 0.0_rt;
122+
123+
// Calculate z-dependent scaling by current
124+
int tp5d_bins = 129;
125+
amrex::ParmParse pp_algo("algo.space_charge");
126+
pp_algo.queryAddWithParser("gauss_charge_z_bins", tp5d_bins);
127+
128+
// Set parameters for charge deposition
129+
bool const is_unity_particle_weight = false;
130+
bool const GetNumberDensity = true;
131+
132+
// Measure beam size, extract the min, max of particle positions
133+
[[maybe_unused]] auto const [x_min, y_min, t_min, x_max, y_max, t_max] =
134+
pc.MinAndMaxPositions();
135+
136+
int const num_bins = tp5d_bins; // Set resolution
137+
amrex::Real const bin_min = t_min;
138+
amrex::Real const bin_max = t_max;
139+
amrex::Real const bin_size = (bin_max - bin_min) / (num_bins - 1); // number of evaluation points
140+
// Allocate memory for the charge profile
141+
amrex::Gpu::DeviceVector<amrex::Real> charge_distribution(num_bins + 1, 0.0);
142+
// Call charge deposition function
143+
impactx::particles::wakefields::DepositCharge1D(pc, charge_distribution, bin_min, bin_size, is_unity_particle_weight);
144+
145+
// Sum up all partial charge histograms to each MPI process to calculate
146+
// the global charge slope.
147+
amrex::ParallelAllReduce::Sum(
148+
charge_distribution.data(),
149+
charge_distribution.size(),
150+
amrex::ParallelDescriptor::Communicator()
151+
);
152+
153+
// Call charge density derivative function
154+
amrex::Gpu::DeviceVector<amrex::Real> slopes(charge_distribution.size() - 1, 0.0);
155+
impactx::particles::wakefields::DerivativeCharge1D(charge_distribution, slopes, bin_size,GetNumberDensity);
156+
157+
amrex::Real const * const beam_profile_slope = slopes.data();
158+
amrex::Real const * const beam_profile = charge_distribution.data();
159+
160+
// group together constants for the momentum push
161+
amrex::ParticleReal const chargesign = charge / std::abs(charge);
162+
163+
amrex::ParallelFor(np, [=] AMREX_GPU_DEVICE (int i) {
164+
// access SoA Real data
165+
amrex::ParticleReal & AMREX_RESTRICT x = part_x[i];
166+
amrex::ParticleReal & AMREX_RESTRICT y = part_y[i];
167+
amrex::ParticleReal z = 0.0_prt; // flatten 3rd dimension
168+
amrex::ParticleReal & AMREX_RESTRICT px = part_px[i];
169+
amrex::ParticleReal & AMREX_RESTRICT py = part_py[i];
170+
amrex::ParticleReal & AMREX_RESTRICT pz = part_pz[i];
171+
172+
// force gather
173+
amrex::GpuArray<amrex::Real, 3> const field_interp =
174+
ablastr::particles::doGatherVectorFieldNodal<2>(
175+
x, y, z,
176+
scf_arr_x, scf_arr_y, scf_arr_z,
177+
invdr,
178+
prob_lo_2D
179+
);
180+
181+
// Update momentae with the 2.5D SC force
182+
int const idx = static_cast<int>((z - bin_min) / bin_size); // Find index position along z
183+
#if (defined(AMREX_DEBUG) || defined(DEBUG)) && !defined(AMREX_USE_GPU)
184+
if (idx < 0 || idx >= num_bins)
185+
{
186+
std::cerr << "Warning: Index out of range for 2.5D SC: " << idx << std::endl;
187+
}
188+
#endif
189+
[[maybe_unused]] amrex::ParticleReal const Fxy = beam_profile[idx] * chargesign;
190+
[[maybe_unused]] amrex::ParticleReal const Fz = beam_profile_slope[idx] * charge;
191+
192+
// push momentum
193+
px += field_interp[0] * Fxy * push_consts;
194+
py += field_interp[1] * Fxy * push_consts;
195+
pz += 0.0_rt;
196+
//pz -= (eintz + pz_push_const) * Fz * push_consts;
197+
198+
// push position is done in the lattice elements
199+
});
200+
}
165201
if (space_charge == SpaceChargeAlgo::True_3D) {
166202
amrex::ParallelFor(np, [=] AMREX_GPU_DEVICE (int i) {
167203
// access SoA Real data

0 commit comments

Comments
 (0)