Skip to content

Commit 5ef6bce

Browse files
AlexanderSinnAlexanderSinn
authored
Use ParticleTileData in plasma ionization (#1315)
1 parent c85a4a0 commit 5ef6bce

1 file changed

Lines changed: 67 additions & 86 deletions

File tree

‎src/particles/plasma/PlasmaParticleContainer.cpp‎

Lines changed: 67 additions & 86 deletions
Original file line numberDiff line numberDiff line change
@@ -399,22 +399,14 @@ IonizationModule (const int lev,
399399
mfi_ion.index(), mfi_ion.LocalTileIndex());
400400
auto& ptile_ion = plevel_ion.at(index);
401401

402-
auto& soa_ion = ptile_ion.GetStructOfArrays(); // For momenta and weights
403-
404402
// Calculation of E0 in SI units for denormalization
405403
const amrex::Real wp = std::sqrt(static_cast<double>(background_density_SI) *
406404
PhysConstSI::q_e*PhysConstSI::q_e /
407405
(PhysConstSI::ep0 * PhysConstSI::m_e) );
408406
const amrex::Real E0 = Hipace::m_normalized_units ?
409407
wp * PhysConstSI::m_e * PhysConstSI::c / PhysConstSI::q_e : 1;
410408

411-
int * const ion_lev = soa_ion.GetIntData(PlasmaIdx::ion_lev).data();
412-
const amrex::Real * const x_prev = soa_ion.GetRealData(PlasmaIdx::x_prev).data();
413-
const amrex::Real * const y_prev = soa_ion.GetRealData(PlasmaIdx::y_prev).data();
414-
const amrex::Real * const uxp = soa_ion.GetRealData(PlasmaIdx::ux_half_step).data();
415-
const amrex::Real * const uyp = soa_ion.GetRealData(PlasmaIdx::uy_half_step).data();
416-
const amrex::Real * const psip =soa_ion.GetRealData(PlasmaIdx::psi_half_step).data();
417-
const auto * idcpup = soa_ion.GetIdCPUData().data();
409+
auto ptd_ion = ptile_ion.getParticleTileData();
418410

419411
// Make Ion Mask and load ADK prefactors
420412
// Ion Mask is necessary to only resize electron particle tile once
@@ -446,12 +438,11 @@ IonizationModule (const int lev,
446438
[=] AMREX_GPU_DEVICE (long ip, const amrex::RandomEngine& engine,
447439
auto depos_order_xy) {
448440

449-
if (amrex::ConstParticleIDWrapper(idcpup[ip]) < 0 ||
450-
amrex::ConstParticleCPUWrapper(idcpup[ip]) != lev) return;
441+
if (!ptd_ion.id(ip).is_valid() || ptd_ion.cpu(ip) != lev) return;
451442

452443
// Avoid temp slice
453-
const amrex::Real xp = x_prev[ip];
454-
const amrex::Real yp = y_prev[ip];
444+
const amrex::Real xp = ptd_ion.rdata(PlasmaIdx::x_prev)[ip];
445+
const amrex::Real yp = ptd_ion.rdata(PlasmaIdx::y_prev)[ip];
455446

456447
// Define field at particle position reals
457448
amrex::ParticleReal ExmByp = 0., EypBxp = 0., Ezp = 0.;
@@ -466,21 +457,25 @@ IonizationModule (const int lev,
466457
const amrex::ParticleReal Eyp = EypBxp - Bxp * phys_const.c;
467458
const amrex::ParticleReal Ep = std::sqrt( Exp*Exp + Eyp*Eyp + Ezp*Ezp )*E0;
468459

460+
const amrex::Real ux = ptd_ion.rdata(PlasmaIdx::ux_half_step)[ip];
461+
const amrex::Real uy = ptd_ion.rdata(PlasmaIdx::uy_half_step)[ip];
462+
const amrex::Real psi = ptd_ion.rdata(PlasmaIdx::psi_half_step)[ip];
463+
469464
// Compute probability of ionization p
470-
const amrex::Real gammap = (1.0_rt + uxp[ip] * uxp[ip]
471-
+ uyp[ip] * uyp[ip]
472-
+ psip[ip]* psip[ip] ) / ( 2.0_rt * psip[ip] );
473-
const int ion_lev_loc = ion_lev[ip];
465+
const amrex::Real gammap = (1.0_rt + ux * ux
466+
+ uy * uy
467+
+ psi * psi ) / ( 2.0_rt * psi );
468+
const int ion_lev_loc = ptd_ion.idata(PlasmaIdx::ion_lev)[ip];
474469
// gamma / (psi + 1) to complete dt for QSA
475-
amrex::Real w_dtau = gammap / psip[ip] * adk_prefactor[ion_lev_loc] *
470+
amrex::Real w_dtau = gammap / psi * adk_prefactor[ion_lev_loc] *
476471
std::pow(Ep, adk_power[ion_lev_loc]) *
477472
std::exp( adk_exp_prefactor[ion_lev_loc]/Ep );
478473
amrex::Real p = 1._rt - std::exp( - w_dtau );
479474

480475
amrex::Real random_draw = amrex::Random(engine);
481476
if (random_draw < p)
482477
{
483-
ion_lev[ip] += 1;
478+
ptd_ion.idata(PlasmaIdx::ion_lev)[ip] += 1;
484479
p_ion_mask[ip] = 1;
485480
amrex::Gpu::Atomic::Add( p_num_new_electrons, 1u ); // ensures thread-safe access when incrementing `p_ip_elec`
486481
}
@@ -500,11 +495,8 @@ IonizationModule (const int lev,
500495
const auto new_size = old_size + num_new_electrons.dataValue();
501496
ptile_elec.resize(new_size);
502497

503-
// Load electron soa and aos after resize
504-
auto arrdata_ion = ptile_ion.GetStructOfArrays().realarray();
505-
auto arrdata_elec = ptile_elec.GetStructOfArrays().realarray();
506-
auto int_arrdata_elec = ptile_elec.GetStructOfArrays().intarray();
507-
auto idcpu_elec = ptile_elec.GetStructOfArrays().GetIdCPUData().data();
498+
// Load electron after resize
499+
auto ptd_elec = ptile_elec.getParticleTileData();
508500

509501
const int init_ion_lev = m_product_pc->m_init_ion_lev;
510502

@@ -521,30 +513,30 @@ IonizationModule (const int lev,
521513

522514
// Copy ion data to new electron
523515
// Set the ionized electron ID to 2 (valid/invalid) for the ionized electrons
524-
amrex::ParticleIDWrapper{idcpu_elec[pidx]} = 2;
525-
amrex::ParticleCPUWrapper{idcpu_elec[pidx]} = lev; // current level
526-
arrdata_elec[PlasmaIdx::x ][pidx] = arrdata_ion[PlasmaIdx::x ][ip];
527-
arrdata_elec[PlasmaIdx::y ][pidx] = arrdata_ion[PlasmaIdx::y ][ip];
528-
529-
arrdata_elec[PlasmaIdx::w ][pidx] = arrdata_ion[PlasmaIdx::w ][ip];
530-
arrdata_elec[PlasmaIdx::ux ][pidx] = 0._rt;
531-
arrdata_elec[PlasmaIdx::uy ][pidx] = 0._rt;
516+
ptd_elec.id(pidx) = 2;
517+
ptd_elec.cpu(pidx) = lev; // current level
518+
ptd_elec.rdata(PlasmaIdx::x )[pidx] = ptd_ion.rdata(PlasmaIdx::x)[ip];
519+
ptd_elec.rdata(PlasmaIdx::y )[pidx] = ptd_ion.rdata(PlasmaIdx::y)[ip];
520+
521+
ptd_elec.rdata(PlasmaIdx::w )[pidx] = ptd_ion.rdata(PlasmaIdx::w)[ip];
522+
ptd_elec.rdata(PlasmaIdx::ux )[pidx] = 0._rt;
523+
ptd_elec.rdata(PlasmaIdx::uy )[pidx] = 0._rt;
532524
// Later we could consider adding a finite temperature to the ionized electrons
533-
arrdata_elec[PlasmaIdx::psi ][pidx] = 1._rt;
534-
arrdata_elec[PlasmaIdx::x_prev ][pidx] = arrdata_ion[PlasmaIdx::x_prev][ip];
535-
arrdata_elec[PlasmaIdx::y_prev ][pidx] = arrdata_ion[PlasmaIdx::y_prev][ip];
536-
arrdata_elec[PlasmaIdx::ux_half_step ][pidx] = 0._rt;
537-
arrdata_elec[PlasmaIdx::uy_half_step ][pidx] = 0._rt;
538-
arrdata_elec[PlasmaIdx::psi_half_step][pidx] = 1._rt;
525+
ptd_elec.rdata(PlasmaIdx::psi )[pidx] = 1._rt;
526+
ptd_elec.rdata(PlasmaIdx::x_prev )[pidx] = ptd_ion.rdata(PlasmaIdx::x_prev)[ip];
527+
ptd_elec.rdata(PlasmaIdx::y_prev )[pidx] = ptd_ion.rdata(PlasmaIdx::y_prev)[ip];
528+
ptd_elec.rdata(PlasmaIdx::ux_half_step )[pidx] = 0._rt;
529+
ptd_elec.rdata(PlasmaIdx::uy_half_step )[pidx] = 0._rt;
530+
ptd_elec.rdata(PlasmaIdx::psi_half_step)[pidx] = 1._rt;
539531
#ifdef HIPACE_USE_AB5_PUSH
540532
#ifdef AMREX_USE_GPU
541533
#pragma unroll
542534
#endif
543535
for (int iforce = PlasmaIdx::Fx1; iforce <= PlasmaIdx::Fpsi5; ++iforce) {
544-
arrdata_elec[iforce][pidx] = 0._rt;
536+
ptd_elec.rdata(iforce)[pidx] = 0._rt;
545537
}
546538
#endif
547-
int_arrdata_elec[PlasmaIdx::ion_lev][pidx] = init_ion_lev;
539+
ptd_elec.idata(PlasmaIdx::ion_lev)[pidx] = init_ion_lev;
548540
}
549541
});
550542

@@ -595,8 +587,6 @@ LaserIonization (const int islice,
595587
mfi_ion.index(), mfi_ion.LocalTileIndex());
596588
auto& ptile_ion = plevel_ion.at(index);
597589

598-
auto& soa_ion = ptile_ion.GetStructOfArrays(); // for momenta and weights
599-
600590
// Calcuation of E0 in SI units for denormalization
601591
const amrex::Real wp = std::sqrt(static_cast<double>(background_density_SI) *
602592
PhysConstSI::q_e*PhysConstSI::q_e /
@@ -607,13 +597,7 @@ LaserIonization (const int islice,
607597
const amrex::Real omega0 = 2.0 * MathConst::pi * phys_const.c / lambda0;
608598
const bool linear_polarization = laser.LinearPolarization();
609599

610-
int * const ion_lev = soa_ion.GetIntData(PlasmaIdx::ion_lev).data();
611-
const amrex::Real * const x_prev = soa_ion.GetRealData(PlasmaIdx::x_prev).data();
612-
const amrex::Real * const y_prev = soa_ion.GetRealData(PlasmaIdx::y_prev).data();
613-
const amrex::Real * const uxp = soa_ion.GetRealData(PlasmaIdx::ux_half_step).data();
614-
const amrex::Real * const uyp = soa_ion.GetRealData(PlasmaIdx::uy_half_step).data();
615-
const amrex::Real * const psip =soa_ion.GetRealData(PlasmaIdx::psi_half_step).data();
616-
const auto * idcpup = soa_ion.GetIdCPUData().data();
600+
auto ptd_ion = ptile_ion.getParticleTileData();
617601

618602
// Make Ion Mask and load ADK prefactors
619603
// Ion Mask is necessary to only resize electron particle tile once
@@ -648,11 +632,10 @@ LaserIonization (const int islice,
648632
auto depos_order_xy) {
649633

650634
// Avoid temp slice
651-
const amrex::Real xp = x_prev[ip];
652-
const amrex::Real yp = y_prev[ip];
635+
const amrex::Real xp = ptd_ion.rdata(PlasmaIdx::x_prev)[ip];
636+
const amrex::Real yp = ptd_ion.rdata(PlasmaIdx::y_prev)[ip];
653637

654-
if (amrex::ConstParticleIDWrapper(idcpup[ip]) < 0 ||
655-
!laser_bounds.contains(xp, yp)) return;
638+
if (!ptd_ion.id(ip).is_valid() || !laser_bounds.contains(xp, yp)) return;
656639

657640
Complex A = 0;
658641
Complex A_dx = 0;
@@ -669,13 +652,17 @@ LaserIonization (const int islice,
669652
amrex::Real Ep = std::sqrt( amrex::abs(Et*Et) + amrex::abs(El*El) );
670653
Ep *= phys_const.m_e * phys_const.c / phys_const.q_e * E0;
671654

655+
const amrex::Real ux = ptd_ion.rdata(PlasmaIdx::ux_half_step)[ip];
656+
const amrex::Real uy = ptd_ion.rdata(PlasmaIdx::uy_half_step)[ip];
657+
const amrex::Real psi = ptd_ion.rdata(PlasmaIdx::psi_half_step)[ip];
658+
672659
// Compute probability of ionization p
673-
const amrex::Real gammap = (1.0_rt + uxp[ip] * uxp[ip]
674-
+ uyp[ip] * uyp[ip]
675-
+ psip[ip]* psip[ip] ) / ( 2.0_rt * psip[ip] );
676-
const int ion_lev_loc = ion_lev[ip];
660+
const amrex::Real gammap = (1.0_rt + ux * ux
661+
+ uy * uy
662+
+ psi * psi ) / ( 2.0_rt * psi );
663+
const int ion_lev_loc = ptd_ion.idata(PlasmaIdx::ion_lev)[ip];
677664
// gamma / (psi + 1) to complete dt for QSA
678-
amrex::Real w_dtau_dc = gammap / psip[ip] * adk_prefactor[ion_lev_loc] *
665+
amrex::Real w_dtau_dc = gammap / psi * adk_prefactor[ion_lev_loc] *
679666
std::pow(Ep, adk_power[ion_lev_loc]) *
680667
std::exp( adk_exp_prefactor[ion_lev_loc]/Ep );
681668

@@ -687,7 +674,7 @@ LaserIonization (const int islice,
687674
amrex::Real random_draw = amrex::Random(engine);
688675
if (random_draw < p)
689676
{
690-
ion_lev[ip] += 1;
677+
ptd_ion.idata(PlasmaIdx::ion_lev)[ip] += 1;
691678
p_ion_mask[ip] = 1;
692679
amrex::Gpu::Atomic::Add( p_num_new_electrons, 1u ); // ensures thread-safe access when incrementing `p_ip_elec`
693680
}
@@ -707,12 +694,8 @@ LaserIonization (const int islice,
707694
const auto new_size = old_size + num_new_electrons.dataValue();
708695
ptile_elec.resize(new_size);
709696

710-
// Load electron soa and aos after resize
711-
auto arrdata_ion = ptile_ion.GetStructOfArrays().realarray();
712-
auto arrdata_elec = ptile_elec.GetStructOfArrays().realarray();
713-
auto int_arrdata_elec = ptile_elec.GetStructOfArrays().intarray();
714-
auto idcpu_elec = ptile_elec.GetStructOfArrays().GetIdCPUData().data();
715-
auto idcpu_ion = ptile_ion.GetStructOfArrays().GetIdCPUData().data();
697+
// Load electron after resize
698+
auto ptd_elec = ptile_elec.getParticleTileData();
716699

717700
const int init_ion_lev = m_product_pc->m_init_ion_lev;
718701

@@ -739,11 +722,10 @@ LaserIonization (const int islice,
739722
if(p_ion_mask[ip] != 0) {
740723

741724
// Avoid temp slice
742-
const amrex::Real xp = x_prev[ip];
743-
const amrex::Real yp = y_prev[ip];
725+
const amrex::Real xp = ptd_ion.rdata(PlasmaIdx::x_prev)[ip];
726+
const amrex::Real yp = ptd_ion.rdata(PlasmaIdx::y_prev)[ip];
744727

745-
if (amrex::ConstParticleIDWrapper(idcpup[ip]) < 0 ||
746-
!laser_bounds.contains(xp, yp)) return;
728+
if (!ptd_ion.id(ip).is_valid() || !laser_bounds.contains(xp, yp)) return;
747729

748730
Complex A = 0;
749731
Complex A_dx = 0;
@@ -759,7 +741,7 @@ LaserIonization (const int islice,
759741
if (linear_polarization) {
760742
// Get the level from which the electron was ionized.
761743
// The -1 is needed as this variable was incremented in the ionization kernel above.
762-
const int ion_lev_loc = ion_lev[ip]-1;
744+
const int ion_lev_loc = ptd_ion.idata(PlasmaIdx::ion_lev)[ip]-1;
763745
const Complex Et = I * A * omega0 + A_dzeta * phys_const.c; // transverse component
764746
const Complex El = - A_dx * phys_const.c; // longitudinal component
765747
amrex::Real Ep = std::sqrt( amrex::abs(Et*Et) + amrex::abs(El*El) );
@@ -791,31 +773,30 @@ LaserIonization (const int islice,
791773
const long pidx = pid + old_size;
792774
// Copy ion data to new electron
793775
// Set the ionized electron ID to 2 (valid/invalid) for the ionized electrons
794-
amrex::ParticleIDWrapper{idcpu_elec[pidx]} = 2;
795-
amrex::ParticleCPUWrapper{idcpu_elec[pidx]} =
796-
amrex::ParticleCPUWrapper{idcpu_ion[pidx]}; // current level
797-
arrdata_elec[PlasmaIdx::x ][pidx] = arrdata_ion[PlasmaIdx::x ][ip];
798-
arrdata_elec[PlasmaIdx::y ][pidx] = arrdata_ion[PlasmaIdx::y ][ip];
799-
arrdata_elec[PlasmaIdx::w ][pidx] = arrdata_ion[PlasmaIdx::w ][ip];
800-
arrdata_elec[PlasmaIdx::ux ][pidx] = ux;
801-
arrdata_elec[PlasmaIdx::uy ][pidx] = uy;
802-
arrdata_elec[PlasmaIdx::psi ][pidx] = std::sqrt(1._rt + ux*ux + uy*uy + uz*uz
776+
ptd_elec.id(pidx) = 2;
777+
ptd_elec.cpu(pidx) = ptd_ion.cpu(ip); // current level
778+
ptd_elec.rdata(PlasmaIdx::x )[pidx] = ptd_ion.rdata(PlasmaIdx::x)[ip];
779+
ptd_elec.rdata(PlasmaIdx::y )[pidx] = ptd_ion.rdata(PlasmaIdx::y)[ip];
780+
ptd_elec.rdata(PlasmaIdx::w )[pidx] = ptd_ion.rdata(PlasmaIdx::w)[ip];
781+
ptd_elec.rdata(PlasmaIdx::ux )[pidx] = ux;
782+
ptd_elec.rdata(PlasmaIdx::uy )[pidx] = uy;
783+
ptd_elec.rdata(PlasmaIdx::psi )[pidx] = std::sqrt(1._rt + ux*ux + uy*uy + uz*uz
803784
+ 0.5_rt*amrex::abs(A*A))-uz;
804-
arrdata_elec[PlasmaIdx::x_prev ][pidx] = arrdata_ion[PlasmaIdx::x_prev][ip];
805-
arrdata_elec[PlasmaIdx::y_prev ][pidx] = arrdata_ion[PlasmaIdx::y_prev][ip];
806-
arrdata_elec[PlasmaIdx::ux_half_step ][pidx] = ux;
807-
arrdata_elec[PlasmaIdx::uy_half_step ][pidx] = uy;
808-
arrdata_elec[PlasmaIdx::psi_half_step][pidx] = std::sqrt(1._rt + ux*ux + uy*uy + uz*uz
785+
ptd_elec.rdata(PlasmaIdx::x_prev )[pidx] = ptd_ion.rdata(PlasmaIdx::x_prev)[ip];
786+
ptd_elec.rdata(PlasmaIdx::y_prev )[pidx] = ptd_ion.rdata(PlasmaIdx::y_prev)[ip];
787+
ptd_elec.rdata(PlasmaIdx::ux_half_step )[pidx] = ux;
788+
ptd_elec.rdata(PlasmaIdx::uy_half_step )[pidx] = uy;
789+
ptd_elec.rdata(PlasmaIdx::psi_half_step)[pidx] = std::sqrt(1._rt + ux*ux + uy*uy + uz*uz
809790
+ 0.5_rt*amrex::abs(A*A))-uz;
810791
#ifdef HIPACE_USE_AB5_PUSH
811792
#ifdef AMREX_USE_GPU
812793
#pragma unroll
813794
#endif
814795
for (int iforce = PlasmaIdx::Fx1; iforce <= PlasmaIdx::Fpsi5; ++iforce) {
815-
arrdata_elec[iforce][pidx] = 0._rt;
796+
ptd_elec.rdata(iforce)[pidx] = 0._rt;
816797
}
817798
#endif
818-
int_arrdata_elec[PlasmaIdx::ion_lev][pidx] = init_ion_lev;
799+
ptd_elec.idata(PlasmaIdx::ion_lev)[pidx] = init_ion_lev;
819800
}
820801
});
821802

0 commit comments

Comments
 (0)