def _after_init(self):
"""This function will be called after WarpX has been initialized but
before any other `afterinit` callbacks."""
self.me = self.sim_ext.getMyProc()
self.n_procs = self.sim_ext.getNProcs()
# A level is needed for many things like level number. For now I'm
# statically setting the default level here. I'm not sure of
# pitfalls or how to handle it more generally yet.
self.lev = self.sim_ext.libwarpx_so.warpx_finestLevel()
self.rho_wrappers = [
fields.RhoFPWrapper(self.lev, False),
fields.RhoFPWrapper(self.lev, True)
]
self.phi_wrappers = [
fields.PhiFPWrapper(self.lev, False),
fields.PhiFPWrapper(self.lev, True)
]
# at this point we are committed to either restarting or starting
# fresh; if this is a fresh start we can delete diags if present
if not self.restart and os.path.exists("diags") and self.me == 0:
shutil.rmtree("diags")
self.initialized = True
def _run_solve(self):
"""Function run on every step to perform the required steps to solve
Poisson's equation."""
# get rho from WarpX
self.rho_data = fields.RhoFPWrapper(0, False)[...][:, 0]
# run superLU solver to get phi
self.solve()
# write phi to WarpX
fields.PhiFPWrapper(0, True)[...] = self.phi
def run_sim(self):
self.sim.step(self.max_steps - self.diag_steps)
self.rho_wrapper = fields.RhoFPWrapper(0, False)
callbacks.installafterstep(self._get_rho_ions)
self.sim.step(self.diag_steps)
if self.sim.extension.getMyProc() == 0:
np.save(f'ion_density_case_{self.n+1}.npy', self.ion_density_array)
def _run_solve(self):
"""Function run on every step to perform the required steps to solve
Poisson's equation."""
# get rho from WarpX
if self.rho_wrapper is None:
self.rho_wrapper = fields.RhoFPWrapper(0, True)
self.rho_data = self.rho_wrapper[Ellipsis][:,:,0]
self.solve()
if self.phi_wrapper is None:
self.phi_wrapper = fields.PhiFPWrapper(0, True)
self.phi_wrapper[Ellipsis] = self.phi
def _get_rho_ions(self):
# deposit the ion density in rho_fp
self.sim.extension.depositChargeDensity('he_ions', 0)
rho_data = fields.RhoFPWrapper(0, False)[...][:, 0]
self.ion_density_array += rho_data / constants.q_e / self.diag_steps