Add initialization of pairwise Coulomb collisions (#155)

* refactor of surface flux diagnostic handling before implementing PHistDiag

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* initial commit of `ParticleHistDiag`

* updated changelog and version number

* added plotting functionality specifically for ZPlane assemblies

* save histogram binning details to file as well

* code cleanup

* changes requested during PR review

* add comment about 2d plotting

* added Coulomb collision installation to picmi.py

* added pairwise Coulomb collision initialization

* remove debugging print statement

* further code changes from PR review

* Fix typo

Co-authored-by: Peter Scherpelz <31747262+peterscherpelz@users.noreply.github.com>

* fix docstring

* cleaned up the logging message for Coulomb scattering

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Peter Scherpelz <31747262+peterscherpelz@users.noreply.github.com>

Examples/Tests/collision/PICMI_inputs_2d.py

@@ -0,0 +1,149 @@
+#!/usr/bin/env python3
+#
+# --- Input file for binary Coulomb collision testing. This input scripts
+# --- runs the same test as inputs_2d but via Python.
+
+from pywarpx import picmi
+
+constants = picmi.constants
+
+#################################
+####### GENERAL PARAMETERS ######
+#################################
+
+max_steps = 150
+
+max_grid_size = 8
+max_level = 0
+nx = max_grid_size
+ny = max_grid_size
+
+xmin = 0
+xmax = 4.154046151855669e2
+ymin = xmin
+ymax = xmax
+
+plasma_density = 1e21
+elec_rms_velocity = 0.044237441120300*constants.c
+ion_rms_velocity = 0.006256118919701*constants.c
+number_per_cell = 200
+
+#################################
+############ NUMERICS ###########
+#################################
+serialize_ics = 1
+verbose = 1
+cfl = 1.0
+
+# Order of particle shape factors
+particle_shape = 1
+
+#################################
+############ PLASMA #############
+#################################
+
+elec_dist = picmi.UniformDistribution(
+    density=plasma_density,
+    rms_velocity=[elec_rms_velocity]*3,
+    directed_velocity=[elec_rms_velocity, 0., 0.]
+)
+
+ion_dist = picmi.UniformDistribution(
+    density=plasma_density,
+    rms_velocity=[ion_rms_velocity]*3,
+)
+
+electrons = picmi.Species(
+    particle_type='electron', name='electron',
+    warpx_do_not_deposit=1,
+    initial_distribution=elec_dist,
+)
+ions = picmi.Species(
+    particle_type='H',
+    name='ion', charge='q_e',
+    mass="5*m_e",
+    warpx_do_not_deposit=1,
+    initial_distribution=ion_dist
+)
+
+#################################
+########## COLLISIONS ###########
+#################################
+
+collision1 = picmi.CoulombCollisions(
+    name='collisions1',
+    species=[electrons, ions],
+    CoulombLog=15.9
+)
+collision2 = picmi.CoulombCollisions(
+    name='collisions2',
+    species=[electrons, electrons],
+    CoulombLog=15.9
+)
+collision3 = picmi.CoulombCollisions(
+    name='collisions3',
+    species=[ions, ions],
+    CoulombLog=15.9
+)
+
+#################################
+###### GRID AND SOLVER ##########
+#################################
+
+grid = picmi.Cartesian2DGrid(
+    number_of_cells=[nx, ny],
+    warpx_max_grid_size=max_grid_size,
+    warpx_blocking_factor=max_grid_size,
+    lower_bound=[xmin, ymin],
+    upper_bound=[xmax, ymax],
+    lower_boundary_conditions=['periodic', 'periodic'],
+    upper_boundary_conditions=['periodic', 'periodic'],
+)
+solver = picmi.ElectromagneticSolver(grid=grid, cfl=cfl)
+
+#################################
+######### DIAGNOSTICS ###########
+#################################
+
+field_diag = picmi.FieldDiagnostic(
+    name='diag1',
+    grid=grid,
+    period=10,
+    data_list=[],
+    write_dir='.',
+    warpx_file_prefix='Python_collisionXZ_plt'
+)
+
+#################################
+####### SIMULATION SETUP ########
+#################################
+
+sim = picmi.Simulation(
+    solver=solver,
+    max_steps=max_steps,
+    verbose=verbose,
+    warpx_serialize_ics=serialize_ics,
+    warpx_collisions=[collision1, collision2, collision3]
+)
+
+sim.add_species(
+    electrons,
+    layout=picmi.PseudoRandomLayout(
+        n_macroparticles_per_cell=number_per_cell, grid=grid
+    )
+)
+sim.add_species(
+    ions,
+    layout=picmi.PseudoRandomLayout(
+        n_macroparticles_per_cell=number_per_cell, grid=grid
+    )
+)
+
+sim.add_diagnostic(field_diag)
+
+#################################
+##### SIMULATION EXECUTION ######
+#################################
+
+#sim.write_input_file('PICMI_inputs_2d')
+sim.step(max_steps)

Examples/Tests/collision/analysis_collision_2d.py

@@ -10,14 +10,14 @@
 # using electron-ion temperature relaxation in 2D.
 # Initially, electrons and ions are both in equilibrium
 # (gaussian) distributions, but have different temperatures.
-# Relaxation occurs to bring the two temperatures to be
-# a final same temperature through collisions.
+# Relaxation occurs to bring the two temperatures to the
+# same final temperature through collisions.
 # The code was tested to be valid, more detailed results
 # were used to obtain an exponential fit with
 # coefficients a and b.
 # This automated test compares the results with the fit.
-# Unrelated to the collision module, we also test the plotfile particle filter function in this
-# analysis script.
+# Unrelated to the collision module, we also test the plotfile
+# particle filter function in this analysis script.
 
 # Possible errors:
 # tolerance: 0.001
@@ -70,7 +70,7 @@
     # load file
     ds  = yt.load( fn )
     ad  = ds.all_data()
-    px  = ad['particle_momentum_x'].to_ndarray()
+    px  = ad[('all', 'particle_momentum_x')].to_ndarray()
     # get time index j
     j = int(fn[-5:])
     # compute error
@@ -87,6 +87,10 @@
 print('tolerance = ', tolerance)
 assert(error < tolerance)
 
+# The second part of the analysis is not done for the Python test
+# since the particle filter function is not accessible from PICMI yet
+if "Python" in last_fn:
+    exit()
 
 ## In the second past of the test, we verify that the diagnostic particle filter function works as
 ## expected. For this, we only use the last simulation timestep.

Python/pywarpx/picmi.py

@@ -859,6 +859,27 @@ def initialize_inputs(self):
         pywarpx.warpx.mirror_z_npoints.append(self.number_of_cells)
 
 
+class CoulombCollisions(picmistandard.base._ClassWithInit):
+    """Custom class to handle setup of binary Coulmb collisions in WarpX. If
+    collision initialization is added to picmistandard this can be changed to
+    inherit that functionality."""
+
+    def __init__(self, name, species, CoulombLog=None, ndt=None, **kw):
+        self.name = name
+        self.species = species
+        self.CoulombLog = CoulombLog
+        self.ndt = ndt
+
+        self.handle_init(kw)
+
+    def initialize_inputs(self):
+        collision = pywarpx.Collisions.newcollision(self.name)
+        collision.type = 'pairwisecoulomb'
+        collision.species = [species.name for species in self.species]
+        collision.CoulombLog = self.CoulombLog
+        collision.ndt = self.ndt
+
+
 class MCCCollisions(picmistandard.base._ClassWithInit):
     """Custom class to handle setup of MCC collisions in WarpX. If collision
     initialization is added to picmistandard this can be changed to inherit

Regression/WarpX-tests.ini

@@ -2053,6 +2053,26 @@ compareParticles = 0
 analysisRoutine = Examples/Tests/collision/analysis_collision_2d.py
 aux1File = Regression/PostProcessingUtils/post_processing_utils.py
 
+[Python_collisionXZ]
+buildDir = .
+inputFile = Examples/Tests/collision/PICMI_inputs_2d.py
+runtime_params =
+customRunCmd = python3 PICMI_inputs_2d.py
+dim = 2
+addToCompileString = USE_PYTHON_MAIN=TRUE
+cmakeSetupOpts = -DWarpX_DIMS=2 -DWarpX_LIB=ON -DWarpX_APP=OFF
+target = pip_install
+restartTest = 0
+useMPI = 1
+numprocs = 1
+useOMP = 1
+numthreads = 1
+compileTest = 0
+doVis = 0
+compareParticles = 0
+analysisRoutine = Examples/Tests/collision/analysis_collision_2d.py
+aux1File = Regression/PostProcessingUtils/post_processing_utils.py
+
 [collisionRZ]
 buildDir = .
 inputFile = Examples/Tests/collision/inputs_rz

mewarpx/changelog.csv

@@ -6,6 +6,8 @@ Version, Physics version, Date,        List of changes
 - Added
   :class:`mewarpx.diags_store.particle_histogram_diagnostic.ParticleHistDiag`
   to track positions of where particles are scraped.
+- Added :class:`mewarpx.coulomb_scattering.PairwiseCoulombScattering` as a
+  wrapper to initialize pairwise Coulomb collisions in simulations.
 
 **Other changes**:
 

mewarpx/mewarpx/coulomb_scattering.py

@@ -4,7 +4,7 @@
 import logging
 
 import numpy as np
-from pywarpx import callbacks
+from pywarpx import callbacks, picmi
 
 from mewarpx.mwxrun import mwxrun
 import mewarpx.utils_store.mwxconstants as constants
@@ -64,13 +64,13 @@ def __init__(self, electron_species, ion_species, log_lambda=None,
 
         print_str = (
             "Initialized electron-ion Coulomb scattering for species "
-            f"{self.collider.name} and {self.field.name}."
+            f"{self.collider.name} and {self.field.name}\n"
         )
         if self.log_lambda is None:
-            print_str += 'The Coulomb logarithm will be calculated.'
+            print_str += "  - the Coulomb logarithm will be calculated"
         else:
             print_str += (
-                'The Coulomb logarithm is fixed at %.2f.' % self.log_lambda
+                f"  - the Coulomb logarithm is fixed at {self.log_lambda:.2f}"
             )
         logger.info(print_str)
 
@@ -217,3 +217,53 @@ def run_scattering_method(self):
             # )
             # print(f"delta energy = {np.sum(E_init - E_final):.3e} eV")
             # assert np.allclose(E_init, E_final)
+
+
+class PairwiseCoulombScattering(object):
+
+    """Wrapper used to initialize pairwise Coulomb collisions."""
+
+    def __init__(self, species1, species2=None, log_lambda=None,
+                 subcycling_steps=None):
+        """
+        Arguments:
+            species1 (:class:`mewarpx.mespecies.Species`): First species that
+                will be scattered.
+            species2 (:class:`mewarpx.mespecies.Species`): Second species that
+                will be scattered. If None species1 will be used.
+            log_lambda (float): If specified, a fixed value for the Coulomb
+                logarithm. If not specified it will be calculated according to
+                the algorithm in Perez et al. (Phys. Plasmas 19, 083104, 2012).
+            subcycling_steps (int): Number of steps between performing pairwise
+                scattering. Default 1 in WarpX code.
+        """
+        self.species1 = species1
+        self.species2 = species2
+        if self.species2 is None:
+            self.species2 = self.species1
+
+        self.log_lambda = log_lambda
+        self.subcycling_steps = subcycling_steps
+
+        print_str = (
+            "Initialized pairwise Coulomb scattering for species "
+            f"{self.species1.name} and {self.species2.name}\n"
+        )
+        if self.log_lambda is None:
+            print_str += "  - the Coulomb logarithm will be calculated"
+        else:
+            print_str += (
+                f"  - the Coulomb logarithm is fixed at {self.log_lambda:.2f}"
+            )
+        logger.info(print_str)
+
+        name = f"coulomb_{self.species1.name}_{self.species2.name}"
+
+        self.collision_object = picmi.CoulombCollisions(
+            name=name, species=[self.species1, self.species2],
+            CoulombLog=self.log_lambda, ndt=self.subcycling_steps
+        )
+
+        if mwxrun.simulation.collisions is None:
+            mwxrun.simulation.collisions = []
+        mwxrun.simulation.collisions.append(self.collision_object)