code: import numpy as np
  import matplotlib.pyplot as plt
  from mpl_toolkits.axes_grid1.axes_divider import make_axes_locatable
  from pywarpx import picmi
  
  # Number of time steps
  max_steps = 100
  
  # Grid
  nx = 128
  nz = 128
  
  # Domain
  xmin = 0.e-6
  zmin = 0.e-6
  xmax = 50.e-6
  zmax = 50.e-6
  
  # Cell size
  dx = (xmax - xmin) / nx
  dz = (zmax - zmin) / nz
  
  # Domain decomposition
  max_grid_size_x = 64
  max_grid_size_z = 64
  
  # PML
  nxpml = 10
  nzpml = 10
  field_boundary = ['open', 'open']
  
  # Spectral order
  nox = 8
  noz = 8
  
  # Guard cells
  nxg = 8
  nzg = 8
  
  # Initialize grid
  grid = picmi.Cartesian2DGrid(number_of_cells = [nx,nz],
                               lower_bound = [xmin,zmin],
                               upper_bound = [xmax,zmax],
                               lower_boundary_conditions = field_boundary,
                               upper_boundary_conditions = field_boundary,
                               guard_cells = [nxg,nzg],
                               moving_window_velocity = [0.,0.,0],
                               warpx_max_grid_size_x = max_grid_size_x,
                               warpx_max_grid_size_y = max_grid_size_z)
  
  # Initialize field solver
  solver = picmi.ElectromagneticSolver(grid=grid, cfl=0.95, method='PSATD',
                                       stencil_order = [nox,noz],
                                       divE_cleaning = 1,
                                       divB_cleaning = 1,
                                       pml_divE_cleaning = 1,
                                       pml_divB_cleaning = 1,
                                       warpx_psatd_update_with_rho = True)
  
  # Initialize diagnostics
  diag_field_list = ["E", "B"]
  field_diag = picmi.FieldDiagnostic(name = 'diag1',
                                     grid = grid,
                                     period = 10,
                                     write_dir = '.',
                                     warpx_file_prefix = 'Python_wrappers_plt',
                                     data_list = diag_field_list)
  
  # Initialize simulation
  sim = picmi.Simulation(solver = solver,
                         max_steps = max_steps,
                         verbose = 1,
                         particle_shape = 'cubic',
                         warpx_current_deposition_algo = 'direct',
                         warpx_particle_pusher_algo = 'boris',
                         warpx_field_gathering_algo = 'energy-conserving',
                         warpx_use_filter = 1)
  
  # Add diagnostics to simulation
  sim.add_diagnostic(field_diag)
  
  # Write input file to run with compiled version
  sim.write_input_file(file_name = 'inputs_2d')
  
  # Whether to include guard cells in data returned by Python wrappers
  include_ghosts = 1
  
  # Compute min and max of fields data
  def compute_minmax(data):
      vmax = np.abs(data).max()
      vmin = -vmax
      return vmin, vmax
  
  # Plot fields data either in valid domain or in PML
  def plot_data(data, pml, title, name):
      fig, ax = plt.subplots(nrows = 1, ncols = 1, gridspec_kw = dict(wspace = 0.5), figsize = [6,5])
      cax = make_axes_locatable(ax).append_axes('right', size='5%', pad='5%')
      lw  = 0.8
      ls = '--'
      if pml:
          # Draw PMLs and ghost regions
          ax.axvline(x = 0       , linewidth = lw, linestyle = ls)
          ax.axvline(x = 0+nxg   , linewidth = lw, linestyle = ls)
          ax.axvline(x = -nxpml  , linewidth = lw, linestyle = ls)
          ax.axvline(x = nx      , linewidth = lw, linestyle = ls)
          ax.axvline(x = nx-nxg  , linewidth = lw, linestyle = ls)
          ax.axvline(x = nx+nxpml, linewidth = lw, linestyle = ls)
          ax.axhline(y = 0       , linewidth = lw, linestyle = ls)
          ax.axhline(y = 0+nzg   , linewidth = lw, linestyle = ls)
          ax.axhline(y = -nzpml  , linewidth = lw, linestyle = ls)
          ax.axhline(y = nz      , linewidth = lw, linestyle = ls)
          ax.axhline(y = nz-nzg  , linewidth = lw, linestyle = ls)
          ax.axhline(y = nz+nzpml, linewidth = lw, linestyle = ls)
          # Annotations
          ax.annotate('PML', xy = (-nxpml//2,nz//2), rotation = 'vertical', ha = 'center', va = 'center')
          ax.annotate('PML', xy = (nx+nxpml//2,nz//2), rotation = 'vertical', ha = 'center', va = 'center')
          ax.annotate('PML', xy = (nx//2,-nzpml//2), rotation = 'horizontal', ha = 'center', va = 'center')
          ax.annotate('PML', xy = (nx//2,nz+nzpml//2), rotation = 'horizontal', ha = 'center', va = 'center')
          ax.annotate('PML ghost', xy = (nxg//2,nz//2), rotation = 'vertical', ha = 'center', va = 'center')
          ax.annotate('PML ghost', xy = (-nxpml-nxg//2,nz//2), rotation = 'vertical', ha = 'center', va = 'center')
          ax.annotate('PML ghost', xy = (nx-nxg//2,nz//2), rotation = 'vertical', ha = 'center', va = 'center')
          ax.annotate('PML ghost', xy = (nx+nxpml+nxg//2,nz//2), rotation = 'vertical', ha = 'center', va = 'center')
          ax.annotate('PML ghost', xy = (nx//2,nzg//2), rotation = 'horizontal', ha = 'center', va = 'center')
          ax.annotate('PML ghost', xy = (nx//2,-nzpml-nzg//2), rotation = 'horizontal', ha = 'center', va = 'center')
          ax.annotate('PML ghost', xy = (nx//2,nz-nzg//2), rotation = 'horizontal', ha = 'center', va = 'center')
          ax.annotate('PML ghost', xy = (nx//2,nz+nzpml+nzg//2), rotation = 'horizontal', ha = 'center', va = 'center')
          # Set extent and sliced data
          extent = np.array([-nxg-nxpml, nx+nxpml+nxg, -nzg-nzpml, nz+nzpml+nzg])
      else:
          # Draw ghost regions
          ax.axvline(x = 0 , linewidth = lw, linestyle = ls)
          ax.axvline(x = nx, linewidth = lw, linestyle = ls)
          ax.axhline(y = 0 , linewidth = lw, linestyle = ls)
          ax.axhline(y = nz, linewidth = lw, linestyle = ls)
          # Annotations
          ax.annotate('ghost', xy = (-nxg//2,nz//2), rotation = 'vertical', ha = 'center', va = 'center')
          ax.annotate('ghost', xy = (nx+nxg//2,nz//2), rotation = 'vertical', ha = 'center', va = 'center')
          ax.annotate('ghost', xy = (nx//2,-nzg//2), rotation = 'horizontal', ha = 'center', va = 'center')
          ax.annotate('ghost', xy = (nx//2,nz+nzg//2), rotation = 'horizontal', ha = 'center', va = 'center')
          # Set extent and sliced data
          extent = np.array([-nxg, nx+nxg, -nzg, nz+nzg])
      X = data[:,:].transpose()
      # Min and max for colorbar
      vmin, vmax = compute_minmax(X)
      # Display data as image
      im = ax.imshow(X = X, origin = 'lower', extent = extent, vmin = vmin, vmax = vmax, cmap = 'seismic')
      # Add colorbar to plot
      fig.colorbar(im, cax = cax)
      # Set label for x- and y-axis, set title
      ax.set_xlabel('x')
      ax.set_ylabel('z')
      ax.set_title(title)
      # Set plot title
      suptitle = 'PML in (x,z), 4 grids 64 x 64'
      plt.suptitle(suptitle)
      # Save figure
      figname = 'figure_' + name + '.png'
      fig.savefig(figname, dpi = 100)
  
  # Initialize fields data (unit pulse) and apply smoothing
  def init_data(data):
      impulse_1d = np.array([1./4., 1./2., 1./4.])
      impulse = np.outer(impulse_1d, impulse_1d)
      data[nx//2-1:nx//2+2,nz//2-1:nz//2+2] = impulse
  
  # Initialize inputs and WarpX instance
  sim.initialize_inputs()
  sim.initialize_warpx()
  
  # Get fields data using Python wrappers
  import pywarpx.fields as pwxf
  Ex = pwxf.ExFPWrapper(include_ghosts = include_ghosts)
  Ey = pwxf.EyFPWrapper(include_ghosts = include_ghosts)
  Ez = pwxf.EzFPWrapper(include_ghosts = include_ghosts)
  Bx = pwxf.BxFPWrapper(include_ghosts = include_ghosts)
  By = pwxf.ByFPWrapper(include_ghosts = include_ghosts)
  Bz = pwxf.BzFPWrapper(include_ghosts = include_ghosts)
  F  = pwxf.FFPWrapper(include_ghosts = include_ghosts)
  G  = pwxf.GFPWrapper(include_ghosts = include_ghosts)
  Expml = pwxf.ExFPPMLWrapper(include_ghosts = include_ghosts)
  Eypml = pwxf.EyFPPMLWrapper(include_ghosts = include_ghosts)
  Ezpml = pwxf.EzFPPMLWrapper(include_ghosts = include_ghosts)
  Bxpml = pwxf.BxFPPMLWrapper(include_ghosts = include_ghosts)
  Bypml = pwxf.ByFPPMLWrapper(include_ghosts = include_ghosts)
  Bzpml = pwxf.BzFPPMLWrapper(include_ghosts = include_ghosts)
  Fpml  = pwxf.FFPPMLWrapper(include_ghosts = include_ghosts)
  Gpml  = pwxf.GFPPMLWrapper(include_ghosts = include_ghosts)
  
  # Initialize fields data in valid domain
  init_data(Ex)
  init_data(Ey)
  init_data(Ez)
  init_data(Bx)
  init_data(By)
  init_data(Bz)
  init_data(F)
  init_data(G)
  
  # Advance simulation until last time step
  sim.step(max_steps)
  
  # Plot E
  plot_data(Ex, pml = False, title = 'Ex', name = 'Ex')
  plot_data(Ey, pml = False, title = 'Ey', name = 'Ey')
  plot_data(Ez, pml = False, title = 'Ez', name = 'Ez')
  
  # Plot B
  plot_data(Bx, pml = False, title = 'Bx', name = 'Bx')
  plot_data(By, pml = False, title = 'By', name = 'By')
  plot_data(Bz, pml = False, title = 'Bz', name = 'Bz')
  
  # F and G
  plot_data(F, pml = False, title = 'F', name = 'F')
  plot_data(G, pml = False, title = 'G', name = 'G')
  
  # Plot E in PML
  plot_data(Expml[:,:,0], pml = True, title = 'Exy in PML', name = 'Exy')
  plot_data(Expml[:,:,1], pml = True, title = 'Exz in PML', name = 'Exz')
  plot_data(Expml[:,:,2], pml = True, title = 'Exx in PML', name = 'Exx')
  plot_data(Eypml[:,:,0], pml = True, title = 'Eyz in PML', name = 'Eyz')
  plot_data(Eypml[:,:,1], pml = True, title = 'Eyx in PML', name = 'Eyx')
  plot_data(Eypml[:,:,2], pml = True, title = 'Eyy in PML', name = 'Eyy') # zero
  plot_data(Ezpml[:,:,0], pml = True, title = 'Ezx in PML', name = 'Ezx')
  plot_data(Ezpml[:,:,1], pml = True, title = 'Ezy in PML', name = 'Ezy') # zero
  plot_data(Ezpml[:,:,2], pml = True, title = 'Ezz in PML', name = 'Ezz')
  
  # Plot B in PML
  plot_data(Bxpml[:,:,0], pml = True, title = 'Bxy in PML', name = 'Bxy')
  plot_data(Bxpml[:,:,1], pml = True, title = 'Bxz in PML', name = 'Bxz')
  plot_data(Bxpml[:,:,2], pml = True, title = 'Bxx in PML', name = 'Bxx')
  plot_data(Bypml[:,:,0], pml = True, title = 'Byz in PML', name = 'Byz')
  plot_data(Bypml[:,:,1], pml = True, title = 'Byx in PML', name = 'Byx')
  plot_data(Bypml[:,:,2], pml = True, title = 'Byy in PML', name = 'Byy') # zero
  plot_data(Bzpml[:,:,0], pml = True, title = 'Bzx in PML', name = 'Bzx')
  plot_data(Bzpml[:,:,1], pml = True, title = 'Bzy in PML', name = 'Bzy') # zero
  plot_data(Bzpml[:,:,2], pml = True, title = 'Bzz in PML', name = 'Bzz')
  
  # Plot F and G in PML
  plot_data(Fpml[:,:,0], pml = True, title = 'Fx in PML', name = 'Fx')
  plot_data(Fpml[:,:,1], pml = True, title = 'Fy in PML', name = 'Fy')
  plot_data(Fpml[:,:,2], pml = True, title = 'Fz in PML', name = 'Fz')
  plot_data(Gpml[:,:,0], pml = True, title = 'Gx in PML', name = 'Gx')
  plot_data(Gpml[:,:,1], pml = True, title = 'Gy in PML', name = 'Gy')
  plot_data(Gpml[:,:,2], pml = True, title = 'Gz in PML', name = 'Gz')
  
  # Check values with benchmarks (precomputed from the same Python arrays)
  def check_values(benchmark, data, rtol, atol):
      passed = np.allclose(benchmark, np.sum(np.abs(data[:,:])), rtol = rtol, atol = atol)
      assert(passed)
  
  rtol = 1e-09
  atol = 1e-12
  
  # E
  check_values(1013263608.6369569, Ex[:,:], rtol, atol)
  check_values(717278253.4505507 , Ey[:,:], rtol, atol)
  check_values(717866566.5718911 , Ez[:,:], rtol, atol)
  # B
  check_values(3.0214509313437636, Bx[:,:], rtol, atol)
  check_values(3.0242765102729985, By[:,:], rtol, atol)
  check_values(3.0214509326970465, Bz[:,:], rtol, atol)
  # F and G
  check_values(3.0188584528062377, F[:,:], rtol, atol)
  check_values(1013672631.8764204, G[:,:], rtol, atol)
  # E in PML
  check_values(364287936.1526477 , Expml[:,:,0], rtol, atol)
  check_values(183582351.3212558 , Expml[:,:,1], rtol, atol)
  check_values(190065766.41491824, Expml[:,:,2], rtol, atol)
  check_values(440581905.9336025 , Eypml[:,:,0], rtol, atol)
  check_values(178117293.6629357 , Eypml[:,:,1], rtol, atol)
  check_values(0.0               , Eypml[:,:,2], rtol, atol)
  check_values(430277101.26568377, Ezpml[:,:,0], rtol, atol)
  check_values(0.0               , Ezpml[:,:,1], rtol, atol)
  check_values(190919663.2167449 , Ezpml[:,:,2], rtol, atol)
  # B in PML
  check_values(1.0565189315366146 , Bxpml[:,:,0], rtol, atol)
  check_values(0.4618191395098556 , Bxpml[:,:,1], rtol, atol)
  check_values(0.6849858273929585 , Bxpml[:,:,2], rtol, atol)
  check_values(1.7228584190213505 , Bypml[:,:,0], rtol, atol)
  check_values(0.47697331996765685, Bypml[:,:,1], rtol, atol)
  check_values(0.0                , Bypml[:,:,2], rtol, atol)
  check_values(1.5183380774611628 , Bzpml[:,:,0], rtol, atol)
  check_values(0.0                , Bzpml[:,:,1], rtol, atol)
  check_values(0.6849858291863835 , Bzpml[:,:,2], rtol, atol)
  # F and G in PML
  check_values(1.7808748509425263, Fpml[:,:,0], rtol, atol)
  check_values(0.0               , Fpml[:,:,1], rtol, atol)
  check_values(0.4307845604625681, Fpml[:,:,2], rtol, atol)
  check_values(536552745.42701197, Gpml[:,:,0], rtol, atol)
  check_values(0.0               , Gpml[:,:,1], rtol, atol)
  check_values(196016270.97767758, Gpml[:,:,2], rtol, atol)
  
apis: numpy.array numpy.abs numpy.outer matplotlib.pyplot.suptitle