diff options
| -rw-r--r-- | Examples/Physics_applications/capacitive_discharge/PICMI_inputs_2d.py | 187 | ||||
| -rwxr-xr-x | Examples/Physics_applications/capacitive_discharge/analysis.py | 18 | ||||
| -rwxr-xr-x | Python/pywarpx/_libwarpx.py | 61 | ||||
| -rw-r--r-- | Python/pywarpx/fields.py | 20 | ||||
| -rw-r--r-- | Regression/WarpX-tests.ini | 2 | ||||
| -rw-r--r-- | Source/Python/WarpXWrappers.cpp | 5 |
6 files changed, 266 insertions, 27 deletions
diff --git a/Examples/Physics_applications/capacitive_discharge/PICMI_inputs_2d.py b/Examples/Physics_applications/capacitive_discharge/PICMI_inputs_2d.py index b597a0f35..b9bc0bb3d 100644 --- a/Examples/Physics_applications/capacitive_discharge/PICMI_inputs_2d.py +++ b/Examples/Physics_applications/capacitive_discharge/PICMI_inputs_2d.py @@ -1,8 +1,16 @@ # --- Input file for MCC testing. There is already a test of the MCC -# --- functionality so this just tests the PICMI interface +# --- functionality. This tests the PICMI interface for the MCC and +# --- provides an example of how an external Poisson solver can be +# --- used for the field solve step. + +# _libwarpx requires the geometry be set before importing +from pywarpx import geometry +geometry.coord_sys = 0 +geometry.prob_lo = [0, 0] import numpy as np -from pywarpx import picmi +from scipy.sparse import csc_matrix, linalg as sla +from pywarpx import picmi, callbacks, _libwarpx, fields constants = picmi.constants @@ -31,12 +39,12 @@ DT = 1.0 / (400 * FREQ) ########################## # --- Number of time steps -max_steps = 10 -diagnostic_intervals = "::5" +max_steps = 50 +diagnostic_intervals = "::50" # --- Grid nx = 128 -ny = 16 +ny = 8 xmin = 0.0 ymin = 0.0 @@ -45,6 +53,162 @@ ymax = D_CA / nx * ny number_per_cell_each_dim = [32, 16] +############################# +# specialized Poisson solver +# using superLU decomposition +############################# + +class PoissonSolverPseudo1D(picmi.ElectrostaticSolver): + + def __init__(self, grid, **kwargs): + """Direct solver for the Poisson equation using superLU. This solver is + useful for pseudo 1D cases i.e. diode simulations with small x extent. + + Arguments: + grid (picmi.Cartesian2DGrid): Instance of the grid on which the + solver will be installed. + """ + super(PoissonSolverPseudo1D, self).__init__( + grid=grid, method=kwargs.pop('method', 'Multigrid'), + required_precision=1, **kwargs + ) + self.rho_wrapper = None + self.phi_wrapper = None + self.time_sum = 0.0 + + def initialize_inputs(self): + """Grab geometrical quantities from the grid. + """ + self.right_voltage = self.grid.potential_xmax + + # set WarpX boundary potentials to None since we will handle it + # ourselves in this solver + self.grid.potential_xmin = None + self.grid.potential_xmax = None + self.grid.potential_ymin = None + self.grid.potential_ymax = None + self.grid.potential_zmin = None + self.grid.potential_zmax = None + + super(PoissonSolverPseudo1D, self).initialize_inputs() + + self.nx = self.grid.nx + self.nz = self.grid.ny + self.dx = (self.grid.xmax - self.grid.xmin) / self.nx + self.dz = (self.grid.ymax - self.grid.ymin) / self.nz + + if not np.isclose(self.dx, self.dz): + raise RuntimeError('Direct solver requires dx = dz.') + + self.nxguardrho = 2 + self.nzguardrho = 2 + self.nxguardphi = 1 + self.nzguardphi = 1 + + self.phi = np.zeros( + (self.nx + 1 + 2*self.nxguardphi, + self.nz + 1 + 2*self.nzguardphi) + ) + + self.decompose_matrix() + + callbacks.installpoissonsolver(self._run_solve) + + def decompose_matrix(self): + """Function to build the superLU object used to solve the linear + system.""" + self.nxsolve = self.nx + 1 + self.nzsolve = self.nz + 3 + + # Set up the computation matrix in order to solve A*phi = rho + A = np.zeros( + (self.nzsolve*self.nxsolve, self.nzsolve*self.nxsolve) + ) + kk = 0 + for ii in range(self.nxsolve): + for jj in range(self.nzsolve): + temp = np.zeros((self.nxsolve, self.nzsolve)) + + if ii == 0 or ii == self.nxsolve - 1: + temp[ii, jj] = 1. + elif ii == 1: + temp[ii, jj] = -2.0 + temp[ii-1, jj] = 1.0 + temp[ii+1, jj] = 1.0 + elif ii == self.nxsolve - 2: + temp[ii, jj] = -2.0 + temp[ii+1, jj] = 1.0 + temp[ii-1, jj] = 1.0 + elif jj == 0: + temp[ii, jj] = 1.0 + temp[ii, -3] = -1.0 + elif jj == self.nzsolve - 1: + temp[ii, jj] = 1.0 + temp[ii, 2] = -1.0 + else: + temp[ii, jj] = -4.0 + temp[ii, jj+1] = 1.0 + temp[ii, jj-1] = 1.0 + temp[ii-1, jj] = 1.0 + temp[ii+1, jj] = 1.0 + + A[kk] = temp.flatten() + kk += 1 + + A = csc_matrix(A, dtype=np.float32) + self.lu = sla.splu(A) + + 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 solve(self): + """The solution step. Includes getting the boundary potentials and + calculating phi from rho.""" + right_voltage = eval( + self.right_voltage, + {'t':_libwarpx.libwarpx.warpx_gett_new(0), 'sin':np.sin, 'pi':np.pi} + ) + left_voltage = 0.0 + + rho = -self.rho_data[ + self.nxguardrho:-self.nxguardrho, self.nzguardrho:-self.nzguardrho + ] / constants.ep0 + + # Construct b vector + nx, nz = np.shape(rho) + source = np.zeros((nx, nz+2), dtype=np.float32) + source[:,1:-1] = rho * self.dx**2 + + source[0] = left_voltage + source[-1] = right_voltage + + # Construct b vector + b = source.flatten() + + flat_phi = self.lu.solve(b) + self.phi[self.nxguardphi:-self.nxguardphi] = ( + flat_phi.reshape(np.shape(source)) + ) + + self.phi[:self.nxguardphi] = left_voltage + self.phi[-self.nxguardphi:] = right_voltage + + # the electrostatic solver in WarpX keeps the ghost cell values as 0 + self.phi[:,:self.nzguardphi] = 0 + self.phi[:,-self.nzguardphi:] = 0 + ########################## # physics components ########################## @@ -136,14 +300,13 @@ grid = picmi.Cartesian2DGrid( bc_ymax = 'periodic', warpx_potential_hi_x = "%.1f*sin(2*pi*%.5e*t)" % (VOLTAGE, FREQ), lower_boundary_conditions_particles=['absorbing', 'periodic'], - upper_boundary_conditions_particles=['absorbing', 'periodic'], - moving_window_velocity = None, - warpx_max_grid_size = nx//4 + upper_boundary_conditions_particles=['absorbing', 'periodic'] ) -solver = picmi.ElectrostaticSolver( - grid=grid, method='Multigrid', required_precision=1e-6 -) +# solver = picmi.ElectrostaticSolver( +# grid=grid, method='Multigrid', required_precision=1e-6 +# ) +solver = PoissonSolverPseudo1D(grid=grid) ########################## # diagnostics @@ -153,7 +316,7 @@ field_diag = picmi.FieldDiagnostic( name = 'diag1', grid = grid, period = diagnostic_intervals, - data_list = ['rho_electrons', 'rho_he_ions', 'phi'], + data_list = ['rho_electrons', 'rho_he_ions'], write_dir = '.', warpx_file_prefix = 'Python_background_mcc_plt' ) diff --git a/Examples/Physics_applications/capacitive_discharge/analysis.py b/Examples/Physics_applications/capacitive_discharge/analysis.py index f37baf912..22d1c725a 100755 --- a/Examples/Physics_applications/capacitive_discharge/analysis.py +++ b/Examples/Physics_applications/capacitive_discharge/analysis.py @@ -2,11 +2,17 @@ # Copyright 2021 Modern Electron -# This script simply checks that the PICMI_inputs_2d.py run output -# diagnostics, which confirms that the PICMI MCC interface works otherwise -# the run would've crashed. +# This script checks that the PICMI_inputs_2d.py run more-or-less matches the +# results from the non-PICMI run. The PICMI run is using an external Poisson +# solver that directly solves the Poisson equation using matrix inversion +# rather than the iterative approach from the MLMG solver. -import glob +import sys +sys.path.append('../../../../warpx/Regression/Checksum/') -files = sorted(glob.glob('Python_background_mcc_plt*'))[1:] -assert len(files) > 0 +import checksumAPI + +my_check = checksumAPI.evaluate_checksum( + 'background_mcc', 'Python_background_mcc_plt00050', + do_particles=True, rtol=7.5e-2 +) diff --git a/Python/pywarpx/_libwarpx.py b/Python/pywarpx/_libwarpx.py index ca87f4f41..eb296478d 100755 --- a/Python/pywarpx/_libwarpx.py +++ b/Python/pywarpx/_libwarpx.py @@ -186,6 +186,8 @@ libwarpx.warpx_getChargeDensityCP.restype = _LP_LP_c_real libwarpx.warpx_getChargeDensityCPLoVects.restype = _LP_c_int libwarpx.warpx_getChargeDensityFP.restype = _LP_LP_c_real libwarpx.warpx_getChargeDensityFPLoVects.restype = _LP_c_int +libwarpx.warpx_getPhiFP.restype = _LP_LP_c_real +libwarpx.warpx_getPhiFPLoVects.restype = _LP_c_int libwarpx.warpx_getParticleBoundaryBufferSize.restype = ctypes.c_int libwarpx.warpx_getParticleBoundaryBuffer.restype = _LP_LP_c_particlereal libwarpx.warpx_getParticleBoundaryBufferScrapedSteps.restype = _LP_LP_c_int @@ -200,6 +202,7 @@ libwarpx.warpx_getJx_nodal_flag.restype = _LP_c_int libwarpx.warpx_getJy_nodal_flag.restype = _LP_c_int libwarpx.warpx_getJz_nodal_flag.restype = _LP_c_int libwarpx.warpx_getRho_nodal_flag.restype = _LP_c_int +libwarpx.warpx_getPhi_nodal_flag.restype = _LP_c_int #libwarpx.warpx_getPMLSigma.restype = _LP_c_real #libwarpx.warpx_getPMLSigmaStar.restype = _LP_c_real @@ -1350,6 +1353,8 @@ def get_mesh_current_density_fp_pml(level, direction, include_ghosts=True): return _get_mesh_field_list(libwarpx.warpx_getCurrentDensityFP_PML, level, direction, include_ghosts) except ValueError: raise Exception('PML not initialized') + + def get_mesh_charge_density_cp(level, include_ghosts=True): ''' @@ -1402,6 +1407,31 @@ def get_mesh_charge_density_fp(level, include_ghosts=True): return _get_mesh_field_list(libwarpx.warpx_getChargeDensityFP, level, None, include_ghosts) +def get_mesh_phi_fp(level, include_ghosts=True): + ''' + + This returns a list of numpy arrays containing the mesh electrostatic + potential data on each grid for this process. This version returns the + potential on the fine patch for the given level. + + The data for the numpy arrays are not copied, but share the underlying + memory buffer with WarpX. The numpy arrays are fully writeable. + + Parameters + ---------- + + level : the AMR level to get the data for + include_ghosts : whether to include ghost zones or not + + Returns + ------- + + A List of numpy arrays. + + ''' + return _get_mesh_field_list(libwarpx.warpx_getPhiFP, level, None, include_ghosts) + + def _get_mesh_array_lovects(level, direction, include_ghosts=True, getlovectsfunc=None): assert(0 <= level and level <= libwarpx.warpx_finestLevel()) @@ -1806,8 +1836,8 @@ def get_mesh_charge_density_cp_lovects(level, include_ghosts=True): def get_mesh_charge_density_fp_lovects(level, include_ghosts=True): ''' - This returns a list of the lo vectors of the arrays containing the mesh electric field - data on each grid for this process. + This returns a list of the lo vectors of the arrays containing the mesh + charge density data on each grid for this process. Parameters ---------- @@ -1824,6 +1854,27 @@ def get_mesh_charge_density_fp_lovects(level, include_ghosts=True): return _get_mesh_array_lovects(level, None, include_ghosts, libwarpx.warpx_getChargeDensityFPLoVects) +def get_mesh_phi_fp_lovects(level, include_ghosts=True): + ''' + + This returns a list of the lo vectors of the arrays containing the mesh + electrostatic potential data on each grid for this process. + + Parameters + ---------- + + level : the AMR level to get the data for + include_ghosts : whether to include ghost zones or not + + Returns + ------- + + A 2d numpy array of the lo vector for each grid with the shape (dims, number of grids) + + ''' + return _get_mesh_array_lovects(level, None, include_ghosts, libwarpx.warpx_getPhiFPLoVects) + + def _get_nodal_flag(getdatafunc): data = getdatafunc() nodal_flag_ref = np.ctypeslib.as_array(data, (dim,)) @@ -1904,3 +1955,9 @@ def get_Rho_nodal_flag(): This returns a 1d array of the nodal flags for Rho along each direction. A 1 means node centered, and 0 cell centered. ''' return _get_nodal_flag(libwarpx.warpx_getRho_nodal_flag) + +def get_Phi_nodal_flag(): + ''' + This returns a 1d array of the nodal flags for Phi along each direction. A 1 means node centered, and 0 cell centered. + ''' + return _get_nodal_flag(libwarpx.warpx_getPhi_nodal_flag) diff --git a/Python/pywarpx/fields.py b/Python/pywarpx/fields.py index f7131f326..4afd42bf8 100644 --- a/Python/pywarpx/fields.py +++ b/Python/pywarpx/fields.py @@ -441,7 +441,7 @@ class _MultiFABWrapper(object): """Sets slices of a decomposed 2D array. """ ix = index[0] - iz = index[2] + iz = index[1] lovects, ngrow = self._getlovects() hivects, ngrow = self._gethivects() @@ -461,7 +461,7 @@ class _MultiFABWrapper(object): ic = None nx = hivects[0,:].max() - ngrow[0] - nz = hivects[2,:].max() - ngrow[1] + nz = hivects[1,:].max() - ngrow[1] # --- Add extra dimensions so that the input has the same number of # --- dimensions as array. @@ -480,7 +480,7 @@ class _MultiFABWrapper(object): ixstop = ix + 1 if isinstance(iz, slice): izstart = max(iz.start or -ngrow[1], -ngrow[1]) - izstop = min(iz.stop or nz + 1 + ngrow[1], nz + self.overlaps[2] + ngrow[1]) + izstop = min(iz.stop or nz + 1 + ngrow[1], nz + self.overlaps[1] + ngrow[1]) else: izstart = iz izstop = iz + 1 @@ -490,13 +490,13 @@ class _MultiFABWrapper(object): # --- The ix1, 2 etc are relative to global indexing ix1 = max(ixstart, lovects[0,i]) ix2 = min(ixstop, lovects[0,i] + fields[i].shape[0]) - iz1 = max(izstart, lovects[2,i]) - iz2 = min(izstop, lovects[2,i] + fields[i].shape[2]) + iz1 = max(izstart, lovects[1,i]) + iz2 = min(izstop, lovects[1,i] + fields[i].shape[1]) if ix1 < ix2 and iz1 < iz2: sss = (slice(ix1 - lovects[0,i], ix2 - lovects[0,i]), - slice(iz1 - lovects[2,i], iz2 - lovects[2,i])) + slice(iz1 - lovects[1,i], iz2 - lovects[1,i])) if ic is not None: sss = tuple(list(sss) + [ic]) @@ -720,6 +720,14 @@ def RhoFPWrapper(level=0, include_ghosts=False): get_fabs=_libwarpx.get_mesh_charge_density_fp, get_nodal_flag=_libwarpx.get_Rho_nodal_flag, level=level, include_ghosts=include_ghosts) + +def PhiFPWrapper(level=0, include_ghosts=False): + return _MultiFABWrapper(direction=None, + get_lovects=_libwarpx.get_mesh_phi_fp_lovects, + get_fabs=_libwarpx.get_mesh_phi_fp, + get_nodal_flag=_libwarpx.get_Phi_nodal_flag, + level=level, include_ghosts=include_ghosts) + def ExCPPMLWrapper(level=1, include_ghosts=False): assert level>0, Exception('Coarse patch only available on levels > 0') return _MultiFABWrapper(direction=0, diff --git a/Regression/WarpX-tests.ini b/Regression/WarpX-tests.ini index 9300b8bd0..792b47edd 100644 --- a/Regression/WarpX-tests.ini +++ b/Regression/WarpX-tests.ini @@ -2323,7 +2323,7 @@ dim = 2 addToCompileString = USE_PYTHON_MAIN=TRUE PYINSTALLOPTIONS="--user --prefix=" restartTest = 0 useMPI = 1 -numprocs = 1 +numprocs = 2 useOMP = 0 numthreads = 0 compileTest = 0 diff --git a/Source/Python/WarpXWrappers.cpp b/Source/Python/WarpXWrappers.cpp index 921dc6b4f..42128c76d 100644 --- a/Source/Python/WarpXWrappers.cpp +++ b/Source/Python/WarpXWrappers.cpp @@ -336,6 +336,7 @@ extern "C" int* warpx_getJy_nodal_flag() {return getFieldNodalFlagData( WarpX::GetInstance().getcurrent(0,1) );} int* warpx_getJz_nodal_flag() {return getFieldNodalFlagData( WarpX::GetInstance().getcurrent(0,2) );} int* warpx_getRho_nodal_flag() {return getFieldNodalFlagData( WarpX::GetInstance().getrho_fp(0) );} + int* warpx_getPhi_nodal_flag() {return getFieldNodalFlagData( WarpX::GetInstance().getphi_fp(0) );} #define WARPX_GET_SCALAR(SCALAR, GETTER) \ amrex::Real** SCALAR(int lev, \ @@ -357,6 +358,10 @@ extern "C" WARPX_GET_LOVECTS_SCALAR(warpx_getChargeDensityCPLoVects, WarpX::GetInstance().getrho_cp) WARPX_GET_LOVECTS_SCALAR(warpx_getChargeDensityFPLoVects, WarpX::GetInstance().getrho_fp) + WARPX_GET_SCALAR(warpx_getPhiFP, WarpX::GetInstance().getphi_fp) + + WARPX_GET_LOVECTS_SCALAR(warpx_getPhiFPLoVects, WarpX::GetInstance().getphi_fp) + #define WARPX_GET_FIELD_PML(FIELD, GETTER) \ amrex::Real** FIELD(int lev, int direction, \ int *return_size, int *ncomps, int **ngrowvect, int **shapes) { \ |