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