Merge branch 'development' into boundary_buffer_structs

Docs/source/usage/parameters.rst

@@ -1824,7 +1824,9 @@ Particle push, charge and current deposition, field gathering
     Available options are: ``direct``, ``esirkepov``, and ``vay``. The default choice
     is ``esirkepov`` for FDTD maxwell solvers but ``direct`` for standard or
     Galilean PSATD solver (i.e. with ``algo.maxwell_solver = psatd``) and
-    for the hybrid-PIC solver (i.e. with ``algo.maxwell_solver = hybrid``).
+    for the hybrid-PIC solver (i.e. with ``algo.maxwell_solver = hybrid``) and for
+    diagnostics output with the electrostatic solvers (i.e., with
+    ``warpx.do_electrostatic = ...``).
     Note that ``vay`` is only available for ``algo.maxwell_solver = psatd``.
 
     1. ``direct``

Docs/source/usage/python.rst

@@ -14,7 +14,7 @@ defining the simulation time, field solver, registered species, etc.
 .. _usage-picmi-parameters:
 
 Classes
-----------
+-------
 
 Simulation and grid setup
 ^^^^^^^^^^^^^^^^^^^^^^^^^
@@ -222,6 +222,8 @@ where ``<n_ranks>`` is the number of MPI ranks used, and ``<python_script>``
 is the name of the script.
 
 
+.. _usage-picmi-extend:
+
 Extending a Simulation from Python
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
@@ -235,85 +237,101 @@ Places in the WarpX loop where callbacks are available include:
 ``afterinit``, ``beforecollisions``, ``aftercollisions``, ``beforeEsolve``, ``afterEsolve``,
 ``beforeInitEsolve``, ``afterInitEsolve``, ``beforedeposition``, ``afterdeposition``,
 ``beforestep``, ``afterstep``, ``afterdiagnostics``,``afterrestart`` and ``oncheckpointsignal``.
-See the examples in *Examples/Tests/ParticleDataPython* for references on how to use
+See the examples in ``Examples/Tests/ParticleDataPython`` for references on how to use
 ``callbacks``.
 
 There are several "hooks" available via the ``libwarpx`` shared library to access and manipulate
 simulation objects (particles, fields and memory buffers) as well as general properties
 (such as processor number). These "hooks" are accessible through the `Simulation.extension` object.
 
-.. autofunction:: pywarpx.picmi.Simulation.extension.getNProcs
+An important object is ``Simulation.extension.warpx``, which is available during simulation run.
+This object is the Python equivalent to the central ``WarpX`` simulation class and provides access to
+field ``MultiFab`` and ``ParticleContainer`` data.
 
-.. autofunction:: pywarpx.picmi.Simulation.extension.getMyProc
+.. function:: pywarpx.picmi.Simulation.extension.warpx.getistep
 
-.. autofunction:: pywarpx.picmi.Simulation.extension.get_nattr
+.. function:: pywarpx.picmi.Simulation.extension.warpx.gett_new
 
-.. autofunction:: pywarpx.picmi.Simulation.extension.get_nattr_species
+.. function:: pywarpx.picmi.Simulation.extension.warpx.evolve
 
-.. autofunction:: pywarpx.picmi.Simulation.extension.getistep
+.. autofunction:: pywarpx.picmi.Simulation.extension.finalize
 
-.. autofunction:: pywarpx.picmi.Simulation.extension.gett_new
+These and other classes are provided through `pyAMReX <https://github.com/AMReX-Codes/pyamrex>`__.
+After the simulation is initialized, pyAMReX can be accessed via
 
-.. autofunction:: pywarpx.picmi.Simulation.extension.evolve
+.. code-block:: python
 
-.. autofunction:: pywarpx.picmi.Simulation.extension.finalize
+   from pywarpx import picmi, libwarpx
+
+   # ... simulation definition ...
+
+   # equivalent to
+   #   import amrex.space3d as amr
+   # for a 3D simulation
+   amr = libwarpx.amr  # picks the right 1d, 2d or 3d variant
+
+.. function:: amr.ParallelDescriptor.NProcs()
 
-.. autofunction:: pywarpx.picmi.Simulation.extension.getProbLo
+.. function:: amr.ParallelDescriptor.MyProc()
 
-.. autofunction:: pywarpx.picmi.Simulation.extension.getProbHi
+.. function:: amr.ParallelDescriptor.IOProcessor()
 
-.. autofunction:: pywarpx.picmi.Simulation.extension.getCellSize
+.. function:: amr.ParallelDescriptor.IOProcessorNumber()
 
 Particles can be added to the simulation at specific positions and with specific
 attribute values:
 
-.. autofunction:: pywarpx.picmi.Simulation.extension.add_particles
+.. code-block:: python
+
+   from pywarpx import particle_containers, picmi
+
+   # ...
+
+   electron_wrapper = particle_containers.ParticleContainerWrapper("electrons")
+
+.. autofunction:: pywarpx.particle_containers.ParticleContainerWrapper.add_particles
 
 Properties of the particles already in the simulation can be obtained with various
 functions.
 
-.. autofunction:: pywarpx.picmi.Simulation.extension.get_particle_count
+.. autofunction:: pywarpx.particle_containers.ParticleContainerWrapper.get_particle_count
 
-.. autofunction:: pywarpx.picmi.Simulation.extension.get_particle_structs
+.. autofunction:: pywarpx.particle_containers.ParticleContainerWrapper.get_particle_structs
 
-.. autofunction:: pywarpx.picmi.Simulation.extension.get_particle_arrays
+.. autofunction:: pywarpx.particle_containers.ParticleContainerWrapper.get_particle_arrays
 
 The ``get_particle_structs()`` and ``get_particle_arrays()`` functions are called
 by several utility functions of the form ``get_particle_{comp_name}`` where
 ``comp_name`` is one of ``x``, ``y``, ``z``, ``r``, ``theta``, ``id``, ``cpu``,
 ``weight``, ``ux``, ``uy`` or ``uz``.
 
-The index of some specific component of the particle data can be obtained.
-
-.. autofunction:: pywarpx.picmi.Simulation.extension.get_particle_comp_index
-
 New components can be added via Python.
 
-.. autofunction:: pywarpx.picmi.Simulation.extension.add_real_comp
+.. autofunction:: pywarpx.particle_containers.ParticleContainerWrapper.add_real_comp
 
 Various diagnostics are also accessible from Python.
 This includes getting the deposited or total charge density from a given species
 as well as accessing the scraped particle buffer. See the example in
-*Examples/Tests/ParticleBoudaryScrape* for a reference on how to interact
+``Examples/Tests/ParticleBoundaryScrape`` for a reference on how to interact
 with scraped particle data.
 
-.. autofunction:: pywarpx.picmi.Simulation.extension.get_species_charge_sum
+.. autofunction:: pywarpx.particle_containers.ParticleContainerWrapper.get_species_charge_sum
 
-.. autofunction:: pywarpx.picmi.Simulation.extension.depositChargeDensity
+.. autofunction:: pywarpx.particle_containers.ParticleContainerWrapper.deposit_charge_density
 
-.. autofunction:: pywarpx.picmi.Simulation.extension.get_particle_boundary_buffer_size
+.. autofunction:: pywarpx.particle_containers.ParticleBoundaryBufferWrapper.get_particle_boundary_buffer_size
 
-.. autofunction:: pywarpx.picmi.Simulation.extension.get_particle_boundary_buffer_structs
+.. autofunction:: pywarpx.particle_containers.ParticleBoundaryBufferWrapper.get_particle_boundary_buffer_structs
 
-.. autofunction:: pywarpx.picmi.Simulation.extension.get_particle_boundary_buffer
+.. autofunction:: pywarpx.particle_containers.ParticleBoundaryBufferWrapper.get_particle_boundary_buffer
 
-.. autofunction:: pywarpx.picmi.Simulation.extension.clearParticleBoundaryBuffer
+.. autofunction:: pywarpx.particle_containers.ParticleBoundaryBufferWrapper.clearParticleBoundaryBuffer
 
 The embedded boundary conditions can be modified when using the electrostatic solver.
 
-.. autofunction:: pywarpx.picmi.Simulation.extension.set_potential_EB
+.. function:: pywarpx.picmi.Simulation.extension.warpx.set_potential_on_eb
 
-Using Python input as a preprocessor
+Using Python Input as a Preprocessor
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
 In this case, only the pure Python version needs to be installed, as described :ref:`here <developers-gnumake-python>`.

Examples/Physics_applications/capacitive_discharge/PICMI_inputs_1d.py

@@ -11,7 +11,7 @@
 from scipy.sparse import csc_matrix
 from scipy.sparse import linalg as sla
 
-from pywarpx import callbacks, fields, particle_containers, picmi
+from pywarpx import callbacks, fields, libwarpx, particle_containers, picmi
 
 constants = picmi.constants
 
@@ -108,7 +108,7 @@ def solve(self):
         calculating phi from rho."""
 
         left_voltage = 0.0
-        t = self.sim_ext.gett_new()
+        t = self.sim.extension.warpx.gett_new(0)
         right_voltage = eval(self.right_voltage)
 
         # Construct b vector
@@ -300,6 +300,7 @@ def setup_run(self):
             warpx_load_balance_intervals=self.max_steps//5000,
             verbose=self.test
         )
+        self.solver.sim = self.sim
 
         self.sim.add_species(
             self.electrons,
@@ -313,7 +314,6 @@ def setup_run(self):
                 n_macroparticle_per_cell=[self.seed_nppc], grid=self.grid
             )
         )
-        self.solver.sim_ext = self.sim.extension
 
         #######################################################################
         # Add diagnostics for the CI test to be happy                         #
@@ -325,6 +325,7 @@ def setup_run(self):
             file_prefix = 'Python_background_mcc_1d_tridiag_plt'
 
         particle_diag = picmi.ParticleDiagnostic(
+            species=[self.electrons, self.ions],
             name='diag1',
             period=0,
             write_dir='.',
@@ -343,7 +344,8 @@ def setup_run(self):
 
     def _get_rho_ions(self):
         # deposit the ion density in rho_fp
-        self.sim.extension.depositChargeDensity('he_ions', 0)
+        he_ions_wrapper = particle_containers.ParticleContainerWrapper('he_ions')
+        he_ions_wrapper.deposit_charge_density(level=0)
 
         rho_data = self.rho_wrapper[...]
         self.ion_density_array += rho_data / constants.q_e / self.diag_steps
@@ -357,7 +359,7 @@ def run_sim(self):
 
         self.sim.step(self.diag_steps)
 
-        if self.sim.extension.getMyProc() == 0:
+        if libwarpx.amr.ParallelDescriptor.MyProc() == 0:
             np.save(f'ion_density_case_{self.n+1}.npy', self.ion_density_array)
 
         # query the particle z-coordinates if this is run during CI testing

Examples/Physics_applications/capacitive_discharge/PICMI_inputs_2d.py

@@ -177,7 +177,7 @@ def solve(self):
         calculating phi from rho."""
         right_voltage = eval(
             self.right_voltage,
-            {'t':sim.extension.gett_new(0), 'sin':np.sin, 'pi':np.pi}
+            {'t': sim.extension.warpx.gett_new(0), 'sin': np.sin, 'pi': np.pi}
         )
         left_voltage = 0.0
 

Examples/Tests/electrostatic_sphere_eb/PICMI_inputs_3d.py

@@ -119,6 +119,6 @@
 
 sim.step(1)
 
-sim.extension.set_potential_EB("2.")
+sim.extension.warpx.set_potential_on_eb("2.")
 
 sim.step(1)

Examples/Tests/langmuir/PICMI_inputs_rz.py

@@ -178,7 +178,7 @@ def calcEz( z, r, k0, w0, wp, t, epsilons) :
     return( Ez_array )
 
 # Current time of the simulation
-t0 = sim.extension.gett_new(0)
+t0 = sim.extension.warpx.gett_new(0)
 
 # Get the raw field data. Note that these are the real and imaginary
 # parts of the fields for each azimuthal mode.

Examples/Tests/ohm_solver_EM_modes/PICMI_inputs.py

@@ -15,7 +15,7 @@
 from mpi4py import MPI as mpi
 import numpy as np
 
-from pywarpx import callbacks, fields, picmi
+from pywarpx import callbacks, fields, libwarpx, picmi
 
 constants = picmi.constants
 
@@ -25,8 +25,6 @@
     warpx_serialize_initial_conditions=True,
     verbose=0
 )
-# make a shorthand for simulation.extension since we use it a lot
-sim_ext = simulation.extension
 
 
 class EMModes(object):
@@ -318,7 +316,7 @@ def _record_average_fields(self):
         similar format as the reduced diagnostic so that the same analysis
         script can be used regardless of the simulation dimension.
         """
-        step = sim_ext.getistep() - 1
+        step = simulation.extension.warpx.getistep(lev=0) - 1
 
         if step % self.diag_steps != 0:
             return
@@ -327,7 +325,7 @@ def _record_average_fields(self):
         By_warpx = fields.BxWrapper()[...]
         Ez_warpx = fields.EzWrapper()[...]
 
-        if sim_ext.getMyProc() != 0:
+        if libwarpx.amr.ParallelDescriptor.MyProc() != 0:
             return
 
         t = step * self.dt

Examples/Tests/ohm_solver_ion_Landau_damping/PICMI_inputs.py

@@ -14,17 +14,16 @@
 from mpi4py import MPI as mpi
 import numpy as np
 
-from pywarpx import callbacks, fields, particle_containers, picmi
+from pywarpx import callbacks, fields, libwarpx, particle_containers, picmi
 
 constants = picmi.constants
 
 comm = mpi.COMM_WORLD
 
 simulation = picmi.Simulation(
     warpx_serialize_initial_conditions=True,
-    verbose=0)
-# make a shorthand for simulation.extension since we use it a lot
-sim_ext = simulation.extension
+    verbose=0
+)
 
 
 class IonLandauDamping(object):
@@ -266,7 +265,8 @@ def setup_run(self):
 
     def text_diag(self):
         """Diagnostic function to print out timing data and particle numbers."""
-        step = sim_ext.getistep(0)
+        step = simulation.extension.warpx.getistep(lev=0) - 1
+
         if step % (self.total_steps // 10) != 0:
             return
 
@@ -290,7 +290,7 @@ def text_diag(self):
             "{step_rate:4.2f} steps/s"
         )
 
-        if sim_ext.getMyProc() == 0:
+        if libwarpx.amr.ParallelDescriptor.MyProc() == 0:
             print(diag_string.format(**status_dict))
 
         self.prev_time = time.time()
@@ -301,14 +301,14 @@ def _record_average_fields(self):
         similar format as the reduced diagnostic so that the same analysis
         script can be used regardless of the simulation dimension.
         """
-        step = sim_ext.getistep() - 1
+        step = simulation.extension.warpx.getistep(lev=0) - 1
 
         if step % self.diag_steps != 0:
             return
 
         Ez_warpx = fields.EzWrapper()[...]
 
-        if sim_ext.getMyProc() != 0:
+        if libwarpx.amr.ParallelDescriptor.MyProc() != 0:
             return
 
         t = step * self.dt