Skip to content

Commit

Permalink
Changed things to fix restarts and make things faster.
Browse files Browse the repository at this point in the history
  • Loading branch information
branmedi committed Sep 10, 2024
1 parent bae11b8 commit 3327402
Show file tree
Hide file tree
Showing 17 changed files with 101 additions and 124 deletions.
20 changes: 7 additions & 13 deletions src/species_advance/species_advance.cc
Original file line number Diff line number Diff line change
Expand Up @@ -30,12 +30,9 @@ checkpt_species( const species_t * sp ) {
CHECKPT_PTR( sp->pb_diag );

#ifdef FIELD_IONIZATION
size_t ionization_energy_size = sp->ionization_energy.size();
CHECKPT_VAL(size_t, ionization_energy_size);
checkpt_data( sp->ionization_energy.data(),
sp->ionization_energy.size() * sizeof(double),
sp->ionization_energy.size() * sizeof(double),
1, 1, 128 );
checkpt_data( sp->ionization_energy,
sp->n_energy * sizeof(double),
sp->n_energy * sizeof(double), 1, 1, 128 );
#endif
}

Expand All @@ -52,11 +49,7 @@ restore_species( void ) {
RESTORE_PTR( sp->pb_diag );

#ifdef FIELD_IONIZATION
size_t ionization_energy_size;
RESTORE_VAL(size_t, ionization_energy_size);

double *ionization_energies = (double *)restore_data();
std::copy(ionization_energies, ionization_energies + ionization_energy_size, sp->ionization_energy.data());
sp->ionization_energy = (double *)restore_data();
#endif

return sp;
Expand Down Expand Up @@ -127,7 +120,7 @@ species( const char * name,
#ifndef FIELD_IONIZATION
float q,
#else
Kokkos::View<double*, Kokkos::HostSpace> ionization_energy,
int n_energy,
float qn,
float qm,
float ql,
Expand Down Expand Up @@ -156,10 +149,11 @@ species( const char * name,
#ifndef FIELD_IONIZATION
sp->q = q;
#else
sp->ionization_energy = ionization_energy;
sp->n_energy = n_energy;
sp->qn = qn;
sp->qm = qm;
sp->ql = ql;
sp->ionization_energy = new double[n_energy];
#endif
sp->m = m;

Expand Down
16 changes: 10 additions & 6 deletions src/species_advance/species_advance.h
Original file line number Diff line number Diff line change
Expand Up @@ -100,10 +100,10 @@ class species_t {
#ifndef FIELD_IONIZATION
short int q; // Species particle charge
#else
Kokkos::View<double*, Kokkos::HostSpace> ionization_energy; // Species ionization energies
float qn; // principal quantum number
float qm; // magnetic quantum number
float ql; // angular momentum quantum number
int n_energy; // Species number of ionization energies
float qn; // principal quantum number
float qm; // magnetic quantum number
float ql; // angular momentum quantum number
#endif
float m; // Species particle rest mass

Expand Down Expand Up @@ -203,6 +203,10 @@ class species_t {
Kokkos::View<int*> clean_up_from;
Kokkos::View<int*> clean_up_to;

#ifdef FIELD_IONIZATION
double * ionization_energy;
#endif

// Init Kokkos Particle Arrays
species_t(int n_particles, int n_pmovers)
{
Expand Down Expand Up @@ -292,7 +296,7 @@ species( const char * name,
#ifndef FIELD_IONIZATION
float q,
#else
Kokkos::View<double*, Kokkos::HostSpace> ionization_energy,
int n_energy,
float qn,
float qm,
float ql,
Expand Down Expand Up @@ -361,7 +365,7 @@ energy_p_kokkos( const species_t * RESTRICT sp,
const interpolator_array_t * RESTRICT ia );

#ifdef FIELD_IONIZATION
Kokkos::View<int*, Kokkos::LayoutLeft>
Kokkos::View<int*, Kokkos::HostSpace>
ionization_states_kokkos( const species_t * RESTRICT sp );

double
Expand Down
18 changes: 9 additions & 9 deletions src/species_advance/standard/advance_p.cc
Original file line number Diff line number Diff line change
Expand Up @@ -949,7 +949,7 @@ advance_p_kokkos_unified_ionize(
ERROR(( "Conversion factors and laser wavelength needs to be passed as grid variables when field ionization is enabled." ));
}

if( (sp!=sp_e || sp->ionization_energy(0)) && n == 0 )
if( (sp!=sp_e || sp->ionization_energy[0]) && n == 0 )
{
ERROR(( "Quantum numbers need to be passed to species struct when field ionization is enabled." ));
}
Expand Down Expand Up @@ -1096,15 +1096,15 @@ advance_p_kokkos_unified_ionize(
haz[LANE] = qdt_2mc*( (fez[LANE] + dx[LANE]*fdezdx[LANE] ) + dy[LANE]*(fdezdy[LANE] + dx[LANE]*fd2ezdxdy[LANE]) );

// ***** Field Ioization *****
if (charge[LANE] != sp->ionization_energy.extent(0)){
if (charge[LANE] != sp->n_energy){
// Declate varviables
bool multiphoton_ionised = false;
float K;

// Check if the particle is fully ionized already
short int N_ionization = abs(charge[LANE]); // Current ionization state of the particle
short int N_ionization_before = N_ionization; // save variable to compare with ionization state after ionization algorithm
short int N_ionization_levels = sp->ionization_energy.extent(0);
short int N_ionization_levels = sp->n_energy;

// code units
float hax_c = (fex[LANE] + dy[LANE]*fdexdy[LANE] ) + dz[LANE]*(fdexdz[LANE] + dy[LANE]*fd2exdydz[LANE]);
Expand Down Expand Up @@ -1132,7 +1132,7 @@ advance_p_kokkos_unified_ionize(
while (ionization_flag == 1 && t_ionize <= dt && N_ionization < N_ionization_levels) {

// Get the appropriate ionization energy
float epsilon_eV = sp->ionization_energy(int(N_ionization)); // [eV], ionization energy
float epsilon_eV = sp->ionization_energy[int(N_ionization)]; // [eV], ionization energy
float epsilon_au = epsilon_eV/27.2; // atomic units, ionization energy

// Calculate stuff
Expand Down Expand Up @@ -2175,20 +2175,20 @@ advance_p_kokkos_gpu_ionize(
Kokkos::View<int> count("count");
Kokkos::deep_copy(count, sp_e->np);

Kokkos::View<double*> epsilon_eV_list_d("epsilon_eV_list_d", sp->ionization_energy.extent(0));
Kokkos::deep_copy(epsilon_eV_list_d, sp->ionization_energy);
Kokkos::View<double*> epsilon_eV_list_d("epsilon_eV_list_d", sp->n_energy);
Kokkos::deep_copy(epsilon_eV_list_d, Kokkos::View<double*, Kokkos::HostSpace>(sp->ionization_energy, sp->n_energy));
float n = sp->qn; // principal quantum number
float m = sp->qm; // magnetic quantum number
float l = sp->ql; // angular momentum quantum number
float lambda_SI = g->lambda;
short int N_ionization_levels = sp->ionization_energy.extent(0);
short int N_ionization_levels = sp->n_energy;

if( t_to_SI == 0 || l_to_SI == 0 || q_to_SI == 0 || m_to_SI == 0 || lambda_SI == 0)
{
ERROR(( "Conversion factors and laser wavelength needs to be passed as grid variables when field ionization is enabled." ));
}

if( (sp!=sp_e || sp->ionization_energy(0) != 0) && n == 0 )
if( (sp!=sp_e || sp->ionization_energy[0] != 0) && n == 0 )
{
ERROR(( "Quantum numbers need to be passed to species struct when field ionization is enabled." ));
}
Expand Down Expand Up @@ -2898,7 +2898,7 @@ advance_p( /**/ species_t * RESTRICT sp,
#endif
KOKKOS_TIC();
#ifdef FIELD_IONIZATION
if (sp != sp_e && sp->ionization_energy(0) != 0){
if (sp != sp_e && sp->ionization_energy[0] != 0){
ADVANCE_P_IONIZE(
sp->k_p_d,
sp->k_p_i_d,
Expand Down
20 changes: 8 additions & 12 deletions src/species_advance/standard/ionization_states.cc
Original file line number Diff line number Diff line change
Expand Up @@ -4,40 +4,36 @@

// This function calculates the number of particles in each ionization state.
#ifdef FIELD_IONIZATION
Kokkos::View<int*, Kokkos::LayoutLeft>
Kokkos::View<int*, Kokkos::HostSpace>
ionization_states_kokkos(const species_t* RESTRICT sp) {

if(!sp) ERROR(("Bad args"));

const long long int np = sp->np;
const k_particles_t& k_particles = sp->k_p_d;

const int N_states = sp->ionization_energy.extent(0)+1; // Include charge state 0
const int N_states = sp->n_energy+1; // Include charge state 0

Kokkos::View<int*, Kokkos::LayoutLeft> charge_counts("charge_counts", N_states);
auto charge_counts_h = Kokkos::create_mirror(charge_counts);
Kokkos::deep_copy(charge_counts_h,charge_counts);
Kokkos::View<int*, Kokkos::HostSpace> charge_counts_h("charge_counts_h", N_states);
Kokkos::View<int*> charge_counts("charge_counts", N_states);
Kokkos::deep_copy(charge_counts,charge_counts_h);

Kokkos::View<int*, Kokkos::LayoutLeft> global_charge_counts("global_charge_counts", N_states);
auto global_charge_counts_h = Kokkos::create_mirror(global_charge_counts);
Kokkos::deep_copy(global_charge_counts_h,global_charge_counts);
Kokkos::View<int*, Kokkos::HostSpace> global_charge_counts_h("global_charge_counts_h", N_states);

// Calculate (local) number of particles in each charge state
Kokkos::parallel_for("ionization_states_kokkos", np, KOKKOS_LAMBDA(const int n) {
long long int charge = k_particles(n, particle_var::charge);
Kokkos::atomic_increment(&charge_counts(charge));
Kokkos::atomic_increment(&charge_counts(charge));
});
Kokkos::fence();
Kokkos::deep_copy(charge_counts_h,charge_counts);


// Perform MPI reductuion over ranks
for (int i = 0; i < N_states; ++i) {
mp_allsum_i( &charge_counts_h(i), &global_charge_counts_h(i), 1 );
};
Kokkos::deep_copy(global_charge_counts,global_charge_counts_h);

// Return the global charge state
return global_charge_counts;
return global_charge_counts_h;
}
#endif
18 changes: 6 additions & 12 deletions src/vpic/dump.cc
Original file line number Diff line number Diff line change
Expand Up @@ -84,9 +84,7 @@ void
vpic_simulation::dump_ionization_states( const char *fname,
int append ) {

Kokkos::View<int*, Kokkos::LayoutLeft> N_ions;
auto N_ions_h = Kokkos::create_mirror(N_ions);
Kokkos::deep_copy(N_ions_h,N_ions);
Kokkos::View<int*, Kokkos::HostSpace> N_ions_h;

species_t *sp;
FileIO fileIO;
Expand All @@ -97,7 +95,7 @@ vpic_simulation::dump_ionization_states( const char *fname,
// Iterate over each species
LIST_FOR_EACH(sp, species_list) {
// Ignore species with ionization energy set to 0, except for electrons
if (std::string(sp->name) != "electron" && sp->ionization_energy(0) == 0) {
if (std::string(sp->name) != "electron" && sp->n_energy == 0) {
// Skip file creation for the "electron" species or when ionization isnt desired for a species
continue;
}
Expand All @@ -116,7 +114,7 @@ vpic_simulation::dump_ionization_states( const char *fname,
if (append == 0) {
// Dynamically generate the layout string
std::string layoutString = "%% Layout\n%% Number of particles in each ionization state\n%% step";
for (size_t i = 0; i <= sp->ionization_energy.extent(0); ++i) {
for (size_t i = 0; i <= sp->n_energy; ++i) {
layoutString += " " + std::to_string(i) + "+";
}
fileIO.print(layoutString.c_str());
Expand All @@ -127,17 +125,13 @@ vpic_simulation::dump_ionization_states( const char *fname,
}
} // if rank


N_ions = ionization_states_kokkos( sp ); // Number of particles in each ionization state

auto N_ions_h = Kokkos::create_mirror(N_ions);
Kokkos::deep_copy(N_ions_h, N_ions);

// Calculate the total number of electrons
// Calculate the total number of electrons and ions
int np_global;
if (std::string(sp->name) == "electron") {
int np = sp->np;
mp_allsum_i( &np, &np_global, 1 );
} else {
N_ions_h = ionization_states_kokkos( sp ); // Number of particles in each ionization state
}

// Print Number of particles in each ionization state
Expand Down
4 changes: 2 additions & 2 deletions src/vpic/vpic.h
Original file line number Diff line number Diff line change
Expand Up @@ -524,7 +524,7 @@ class vpic_simulation {
#if !defined(FIELD_IONIZATION)
double q,
#else
Kokkos::View<double*, Kokkos::HostSpace> ionization_energy,
int n_energy,
double qn, // quantum numbers n,m,l
double qm,
double ql,
Expand All @@ -546,7 +546,7 @@ class vpic_simulation {
#if !defined(FIELD_IONIZATION)
(float)q,
#else
(Kokkos::View<double*,Kokkos::HostSpace>)ionization_energy,
(int) n_energy,
(float)qn,
(float)qm,
(float)ql,
Expand Down
9 changes: 4 additions & 5 deletions test/integrated/legacy/accel.deck
Original file line number Diff line number Diff line change
Expand Up @@ -71,12 +71,11 @@ begin_initialization {
grid->l_to_SI = 1;
grid->q_to_SI = 1;
grid->m_to_SI = 1;
Kokkos::View<double*, Kokkos::HostSpace> ionization_energy("my_kokkos_view", 1);
double ionization_energy_values[] = {0}; // in eV
ionization_energy(0) = ionization_energy_values[0];
this_sp = define_species( "test_species", ionization_energy, 1,0,0, 1., npart, npart, 0, 0 );
this_sp = define_species( "test_species", 1, 1,0,0, 1., npart, npart, 0, 0 );
this_sp->ionization_energy[0] = 0;
// electron needs to be defined when FI is enabled
species_t * electron = define_species("electron", ionization_energy, 0,0,0, 1., npart, npart, 0, 0);
species_t * electron = define_species("electron", 1, 0,0,0, 1., npart, npart, 0, 0);
electron->ionization_energy[0] = 0;
#else
this_sp = define_species( "test_species", 1., 1., npart, npart, 0, 0 );
#endif
Expand Down
9 changes: 4 additions & 5 deletions test/integrated/legacy/cyclo.deck
Original file line number Diff line number Diff line change
Expand Up @@ -47,12 +47,11 @@ begin_initialization {
grid->l_to_SI = 1;
grid->q_to_SI = 1;
grid->m_to_SI = 1;
Kokkos::View<double*, Kokkos::HostSpace> ionization_energy("my_kokkos_view", 1);
double ionization_energy_values[] = {0}; // in eV
ionization_energy(0) = ionization_energy_values[0];
this_sp = define_species( "test_species", ionization_energy, 1,0,0, 1, npart, npart, 0, 0 );
this_sp = define_species( "test_species", 1, 1,0,0, 1, npart, npart, 0, 0 );
this_sp->ionization_energy[0] = 0;
// electron needs to be defined when FI is enabled
species_t * electron = define_species("electron",ionization_energy, 0,0,0, 1., npart, npart, 0, 0);
species_t * electron = define_species("electron",1, 0,0,0, 1., npart, npart, 0, 0);
electron->ionization_energy[0] = 0;
#else
this_sp = define_species( "test_species", 2, 1, npart, npart, 0, 0 );
#endif
Expand Down
9 changes: 4 additions & 5 deletions test/integrated/legacy/inbndj.deck
Original file line number Diff line number Diff line change
Expand Up @@ -39,12 +39,11 @@ begin_initialization {
grid->l_to_SI = 1;
grid->q_to_SI = 1;
grid->m_to_SI = 1;
Kokkos::View<double*, Kokkos::HostSpace> ionization_energy("my_kokkos_view", 1);
double ionization_energy_values[] = {0}; // in eV
ionization_energy(0) = ionization_energy_values[0];
this_sp = define_species( "test_species", ionization_energy, 1,0,0, 1., NPART, NPART, 0, 0 );
this_sp = define_species( "test_species", 1, 1,0,0, 1., NPART, NPART, 0, 0 );
this_sp->ionization_energy[0] = 0;
// electron needs to be defined when FI is enabled
species_t * electron = define_species("electron", ionization_energy, 0,0,0, 1., NPART, NPART, 0, 0);
species_t * electron = define_species("electron", 1, 0,0,0, 1., NPART, NPART, 0, 0);
electron->ionization_energy[0] = 0;
#else
this_sp = define_species( "test_species", 2., 1., NPART, NPART, 0, 0 );
#endif
Expand Down
9 changes: 4 additions & 5 deletions test/integrated/legacy/interpe.deck
Original file line number Diff line number Diff line change
Expand Up @@ -49,12 +49,11 @@ begin_initialization {
grid->l_to_SI = 1;
grid->q_to_SI = 1;
grid->m_to_SI = 1;
Kokkos::View<double*, Kokkos::HostSpace> ionization_energy("my_kokkos_view", 1);
double ionization_energy_values[] = {0}; // in eV
ionization_energy(0) = ionization_energy_values[0];
this_sp = define_species( "test_species", ionization_energy, 1,0,0, 2, npart, npart, 0, 0 );
this_sp = define_species( "test_species", 1, 1,0,0, 2, npart, npart, 0, 0 );
this_sp->ionization_energy[0] = 0;
// electron needs to be defined when FI is enabled
species_t * electron = define_species("electron", ionization_energy, 0,0,0, 1., npart, npart, 0, 0);
species_t * electron = define_species("electron", 1, 0,0,0, 1., npart, npart, 0, 0);
electron->ionization_energy[0] = 0;
#else
this_sp = define_species( "test_species", 2, 2, npart, npart, 0, 0 );
#endif
Expand Down
9 changes: 4 additions & 5 deletions test/integrated/to_completion/dump.deck
Original file line number Diff line number Diff line change
Expand Up @@ -146,11 +146,10 @@ begin_initialization {
grid->l_to_SI = 1;
grid->q_to_SI = 1;
grid->m_to_SI = 1;
Kokkos::View<double*, Kokkos::HostSpace> ionization_energy("my_kokkos_view", 1);
double ionization_energy_values[] = {0}; // in eV
ionization_energy(0) = ionization_energy_values[0];
ion = define_species( "ion", ionization_energy, 1,0,0, mi, 1.5*Ni/nproc(), -1, 40, 1 );
electron = define_species( "electron", ionization_energy, 1,0,0, me, 1.5*Ne/nproc(), -1, 20, 1 );
ion = define_species( "ion", 1, 1,0,0, mi, 1.5*Ni/nproc(), -1, 40, 1 );
electron = define_species( "electron", 1, 1,0,0, me, 1.5*Ne/nproc(), -1, 20, 1 );
ion->ionization_energy[0] = 0;
electron->ionization_energy[0] = 0;
#else
ion = define_species( "ion", ec, mi, 1.5*Ni/nproc(), -1, 40, 1 );
electron = define_species( "electron", -ec, me, 1.5*Ne/nproc(), -1, 20, 1 );
Expand Down
18 changes: 10 additions & 8 deletions test/integrated/to_completion/reconnection_test.deck
Original file line number Diff line number Diff line change
Expand Up @@ -224,15 +224,17 @@ begin_initialization {
grid->l_to_SI = 1;
grid->q_to_SI = 1;
grid->m_to_SI = 1;
Kokkos::View<double*, Kokkos::HostSpace> ionization_energy("my_kokkos_view", 1);
double ionization_energy_values[] = {0}; // in eV
ionization_energy(0) = ionization_energy_values[0];
electron1 = define_species("electron1", ionization_energy, 1,0,0, me, nmax, nmovers, electron_sort_interval, sort_method);
electron2 = define_species("electron2", ionization_energy, 1,0,0, me, nmax, nmovers, electron_sort_interval, sort_method);
ion1 = define_species("ion1", ionization_energy, 1,0,0, mi, nmax, nmovers, ion_sort_interval, sort_method);
ion2 = define_species("ion2", ionization_energy, 1,0,0, mi, nmax, nmovers, ion_sort_interval, sort_method);
electron1 = define_species("electron1", 1, 1,0,0, me, nmax, nmovers, electron_sort_interval, sort_method);
electron2 = define_species("electron2", 1, 1,0,0, me, nmax, nmovers, electron_sort_interval, sort_method);
ion1 = define_species("ion1", 1, 1,0,0, mi, nmax, nmovers, ion_sort_interval, sort_method);
ion2 = define_species("ion2", 1, 1,0,0, mi, nmax, nmovers, ion_sort_interval, sort_method);
ion1->ionization_energy[0] = 0;
ion2->ionization_energy[0] = 0;
electron2->ionization_energy[0] = 0;
electron2->ionization_energy[0] = 0;
// electron needs to be defined when FI is enabled
species_t *electron = define_species("electron", ionization_energy, 0,0,0, me, nmax, nmovers, electron_sort_interval, sort_method);
species_t *electron = define_species("electron", 1, 0,0,0, me, nmax, nmovers, electron_sort_interval, sort_method);
electron->ionization_energy[0] = 0;
#else
electron1 = define_species("electron1",-ec, me, nmax, nmovers, electron_sort_interval, sort_method);
electron2 = define_species("electron2",-ec, me, nmax, nmovers, electron_sort_interval, sort_method);
Expand Down
Loading

0 comments on commit 3327402

Please sign in to comment.