BLAST-WarpX
diff --git a/‎Examples/Physics_applications/spacecraft_charging/PICMI_inputs_rz.py
+299 b/‎Examples/Physics_applications/spacecraft_charging/PICMI_inputs_rz.py
+299
diff --git a/‎Examples/Physics_applications/spacecraft_charging/analysis.py
+78 b/‎Examples/Physics_applications/spacecraft_charging/analysis.py
+78
@@ -0,0 +1,299 @@
+#!/usr/bin/env python3
+#
+# --- Input file for spacecraft charging testing in RZ.
+# --- This input defines a conducting sphere (spacecraft) immersed in a thermal
+# --- plasma with the same given initial conditions as in the article:
+# --- (*) J. Deca, G. Lapenta, R. Marchand, S. Markidis;
+# ---     Spacecraft charging analysis with the implicit particle-in-cell code iPic3D.
+# ---     Part III. A. pages 3-4
+# ---     Phys. Plasmas 1 October 2013; 20 (10): 102902. https://doi.org/10.1063/1.4826951.
+# --- The conducting sphere starts with an initial potential of 1V and will interact with
+# --- the surrounding plasma, initially static. The charging of the spacecraft - by accumulation
+# --- of electrons - leads to a decrease of the potential on the surface over the time
+# --- until reaching an equilibrium floating potential of ~144.5 V (*).
+
+from mpi4py import MPI as mpi
+import numpy as np
+import scipy.constants as scc
+
+from pywarpx import picmi
+from pywarpx.callbacks import installafterEsolve, installafterInitEsolve
+from pywarpx.fields import ExWrapper, EzWrapper, PhiFPWrapper, RhoFPWrapper
+from pywarpx.particle_containers import ParticleBoundaryBufferWrapper
+
+
+# Utilities
+class SpaceChargeFieldCorrector(object):
+    """
+    Class used by the callback functions to calculate the
+    correct charge on the spacecraft at each initialisation.
+    """
+    def __init__(self):
+        self.saved_first_iteration_fields = False
+        self.spacecraft_potential = 1. # Initial voltage: 1V
+        self.spacecraft_capacitance = None
+
+    def correct_space_charge_fields(self, q=None):
+        """
+        Function that will be called at each iteration,
+        after each electrostatic solve in WarpX
+        """
+        assert self.saved_first_iteration_fields
+
+        # Compute the charge that WarpX thinks there is on the spacecraft
+        # from phi and rho after the Poisson solver
+        q_v = compute_virtual_charge_on_spacecraft()
+        if q is None:
+            q = compute_actual_charge_on_spacecraft()
+
+        # Correct fields so as to recover the actual charge
+        Er = ExWrapper(include_ghosts=True)[:,:]
+        Er[...] = Er[...]+(q - q_v)*self.normalized_Er[...]
+        Ez = EzWrapper(include_ghosts=True)[:,:]
+        Ez[...]  += (q - q_v)*self.normalized_Ez[...]
+        phi = PhiFPWrapper(include_ghosts=True)[:,:]
+        phi[...]  += (q - q_v)*self.normalized_phi[...]
+        self.spacecraft_potential += (q - q_v)*self.spacecraft_capacitance
+        sim.extension.warpx.set_potential_on_eb( "%f" %self.spacecraft_potential )
+        print('Setting potential to %f' %self.spacecraft_potential)
+
+        # Confirm that the charge on the spacecraft is now correct
+        compute_virtual_charge_on_spacecraft()
+
+    def save_normalized_vacuum_Efields(self,):
+        # Compute the charge that WarpX thinks there is on the spacecraft
+        # from phi and rho after the Poisson solver
+        q_v = compute_virtual_charge_on_spacecraft()
+        self.spacecraft_capacitance = 1./q_v # the potential was set to 1V
+
+        # Check that this iteration corresponded to a vacuum solve
+        rho = RhoFPWrapper(include_ghosts=False)
+
+        # In principle, we should check that `rho` is exactly 0
+        # However, due to machine precision errors when adding the charge
+        # of ions and electrons, this can be slightly different than 0
+        assert np.all( abs(rho[...]) < 1.e-11 )
+
+        # Record fields
+        Er = ExWrapper(include_ghosts=True)[:,:]
+        self.normalized_Er = Er[...] /q_v
+        Ez = EzWrapper(include_ghosts=True)[:,:]
+        self.normalized_Ez = Ez[...] /q_v
+        phi = PhiFPWrapper(include_ghosts=True)[:,:]
+        self.normalized_phi = phi[...] /q_v
+
+        self.saved_first_iteration_fields = True
+        self.correct_space_charge_fields(q=0)
+
+
+def compute_virtual_charge_on_spacecraft():
+    """
+    Given that we asked WarpX to solve the Poisson
+    equation with phi=1 on the spacecraft and phi=0
+    on the boundary of the domain, compute the charge
+    that WarpX thinks there should be on the spacecraft.
+    """
+    # Get global array for the whole domain (across MPI ranks)
+    phi = PhiFPWrapper(include_ghosts=False)[:,:]
+    rho = RhoFPWrapper(include_ghosts=False)[:,:]
+
+    # Check that this codes correspond to the global size of the box
+    assert phi.shape == (nr+1, nz+1)
+    assert rho.shape == (nr+1, nz+1)
+
+    dr, dz = sim.extension.warpx.Geom(lev=0).data().CellSize()
+
+    # Compute integral of grad phi over surfaces of the domain
+    r = np.linspace(rmin, rmax, len(phi), endpoint=False) + (rmax - rmin) / (2 * len(phi)) #shift of the r points because the derivaties are calculated in the middle
+    face_z0 = 2 * np.pi *  1./dz * ( (phi[:,0]-phi[:,1]) * r ).sum() * dr #here I am assuming that phi is a numpy array that can handle elementwise mult
+    face_zend = 2 * np.pi * 1./dz * ( (phi[:,-1]-phi[:,-2]) * r ).sum() * dr
+    face_rend = 2 * np.pi * 1./dr*((phi[-1,:]-phi[-2,:]) * rmax).sum() * dz
+    grad_phi_integral = face_z0 + face_zend + face_rend
+
+    # Compute integral of rho over volume of the domain
+    # (i.e. total charge of the plasma particles)
+    rho_integral = 0.0
+    for k in range(1, nz-1):
+        for i in range(1, nr-1):
+            rho_integral += rho[i,k] * r[i] * dr * dz
+
+    # Due to an oddity in WarpX (which will probably be solved later)
+    # we need to multiply `rho` by `-epsilon_0` to get the correct charge
+    rho_integral *= 2 * np.pi * -scc.epsilon_0 #does this oddity still exist?
+
+    # Compute charge of the spacecraft, based on Gauss theorem
+    q_spacecraft = - rho_integral - scc.epsilon_0 * grad_phi_integral
+    print('Virtual charge on the spacecraft: %e' %q_spacecraft)
+    return q_spacecraft
+
+
+def compute_actual_charge_on_spacecraft():
+    """
+    Compute the actual charge on the spacecraft,
+    by counting how many electrons and protons
+    were collected by the WarpX embedded boundary (EB)
+    """
+    charge = {'electrons': -scc.e, 'protons': scc.e}
+    q_spacecraft = 0
+    particle_buffer = ParticleBoundaryBufferWrapper()
+    for species in charge.keys():
+        weights = particle_buffer.get_particle_boundary_buffer(species, 'eb', 'w', 0)
+        sum_weights_over_tiles = sum([w.sum() for w in weights])
+
+        # Reduce across all MPI ranks
+        ntot = float(mpi.COMM_WORLD.allreduce(sum_weights_over_tiles, op=mpi.SUM))
+        print('Total number of %s collected on spacecraft: %e'%(species, ntot))
+        q_spacecraft += ntot * charge[species]
+
+    print('Actual charge on the spacecraft: %e' %q_spacecraft)
+    return q_spacecraft
+
+
+##########################
+# numerics parameters
+##########################
+
+dt=1.27e-8
+
+# --- Nb time steps
+max_steps = 1000
+diagnostic_interval = 10
+
+# --- grid
+nr = 40
+nz= 80
+
+rmin = 0.0
+rmax = 3
+zmin = -3
+zmax = 3
+
+number_per_cell =5
+number_per_cell_each_dim = [10,1, 1]
+
+
+##########################
+# physics components
+##########################
+
+n = 7.0e9 #plasma density #particles/m^3
+Te = 85 #Electron temp in eV
+Ti = 0.05 * Te #Ion temp in eV
+qe = picmi.constants.q_e #elementary charge
+m_e = picmi.constants.m_e #electron mass
+m_i = 1836.0 * m_e #mass of ion
+v_eth = (qe * Te / m_e) ** 0.5
+v_pth = (qe * Ti / m_i) ** 0.5
+
+# nothing to change in the distribution function?
+e_dist = picmi.UniformDistribution(density = n, rms_velocity=[v_eth, v_eth, v_eth] )
+e_dist2 = picmi.UniformFluxDistribution(
+    flux=n*v_eth/(2*np.pi)**.5, # Flux for Gaussian with vmean=0
+    surface_flux_position=3,
+    flux_direction=-1, flux_normal_axis='r',
+    gaussian_flux_momentum_distribution=True,
+    rms_velocity=[v_eth, v_eth, v_eth] )
+electrons = picmi.Species(particle_type='electron',
+                          name='electrons',
+                          initial_distribution=[e_dist,e_dist2],
+                          warpx_save_particles_at_eb=1)
+
+p_dist = picmi.UniformDistribution(density = n, rms_velocity=[v_pth, v_pth, v_pth] )
+p_dist2 = picmi.UniformFluxDistribution(
+    flux=n*v_pth/(2*np.pi)**.5, # Flux for Gaussian with vmean=0
+    surface_flux_position=3,
+    flux_direction=-1, flux_normal_axis='r',
+    gaussian_flux_momentum_distribution=True,
+    rms_velocity=[v_pth, v_pth, v_pth] )
+protons = picmi.Species(particle_type='proton',
+                        name='protons',
+                        initial_distribution=[p_dist,p_dist2],
+                        warpx_save_particles_at_eb=1)
+
+
+##########################
+# numerics components
+##########################
+
+grid = picmi.CylindricalGrid(
+    number_of_cells = [nr, nz],
+    n_azimuthal_modes = 1,
+    lower_bound = [rmin, zmin],
+    upper_bound = [rmax, zmax],
+    lower_boundary_conditions = ['none', 'dirichlet'],
+    upper_boundary_conditions =  ['dirichlet', 'dirichlet'],
+    lower_boundary_conditions_particles = ['absorbing', 'reflecting'],
+    upper_boundary_conditions_particles =  ['absorbing', 'reflecting']
+)
+
+solver = picmi.ElectrostaticSolver(
+    grid=grid, method='Multigrid',
+    warpx_absolute_tolerance=1e-7
+)
+
+embedded_boundary = picmi.EmbeddedBoundary(
+    implicit_function="-(x**2+y**2+z**2-radius**2)",
+    potential=1., # arbitrary value ; this will be corrected by a callback function
+    radius = 0.3277
+)
+
+
+##########################
+# diagnostics
+##########################
+
+field_diag = picmi.FieldDiagnostic(
+    name = 'diag1',
+    grid = grid,
+    period = diagnostic_interval,
+    data_list = ['Er', 'Ez', 'phi', 'rho',
+                 'rho_electrons', 'rho_protons'],
+    warpx_format = 'openpmd',
+    write_dir = '.',
+    warpx_file_prefix = 'spacecraft_charging_plt'
+)
+
+part_diag = picmi.ParticleDiagnostic(name = 'diag1',
+    period = diagnostic_interval,
+    species = [electrons, protons],
+    warpx_format = 'openpmd',
+    write_dir = '.',
+    warpx_file_prefix = 'spacecraft_charging_plt'
+)
+
+##########################
+# simulation setup
+##########################
+
+sim = picmi.Simulation(
+    solver = solver,
+    time_step_size = dt,
+    max_steps = max_steps,
+    warpx_embedded_boundary=embedded_boundary,
+    warpx_amrex_the_arena_is_managed=1,
+    warpx_random_seed=1
+)
+
+layout1=picmi.GriddedLayout(n_macroparticle_per_cell=number_per_cell_each_dim,
+                            grid=grid)
+layout2=picmi.PseudoRandomLayout(n_macroparticles_per_cell=number_per_cell,
+                                 grid=grid)
+sim.add_species(electrons,
+                layout = [layout1,layout2])
+
+sim.add_species(protons,
+                layout = [layout1,layout2])
+
+sim.add_diagnostic(field_diag)
+sim.add_diagnostic(part_diag)
+
+##########################
+# simulation run
+##########################
+
+spc = SpaceChargeFieldCorrector()
+
+installafterInitEsolve( spc.save_normalized_vacuum_Efields )
+installafterEsolve( spc.correct_space_charge_fields )
+
+sim.step(max_steps)
@@ -0,0 +1,78 @@
+#!/usr/bin/env python
+
+"""
+This script tests the potential-time profile on the surface of
+a conducting sphere (spacecraft) immersed in an initially static
+thermal plasma. The potential on the spacecraft decreases over
+the time to reach an equilibrium floating potential.
+
+An input Python file PICMI_inputs_rz.py is used.
+
+The test will check the curve fitting parameters v0 and tau defined
+by the following exponential function: phi(t)=v0(1-exp(-t/tau))
+"""
+
+import os
+import sys
+
+import matplotlib.pyplot as plt
+import numpy as np
+from openpmd_viewer import OpenPMDTimeSeries
+from scipy.optimize import curve_fit
+import yt
+
+yt.funcs.mylog.setLevel(0)
+sys.path.insert(1, '../../../../warpx/Regression/Checksum/')
+import checksumAPI
+
+# Open plotfile specified in command line
+filename = sys.argv[1]
+test_name = os.path.split(os.getcwd())[1]
+checksumAPI.evaluate_checksum(test_name, filename, output_format='openpmd')
+
+ts = OpenPMDTimeSeries('./spacecraft_charging_plt')
+dt = 1.27e-8
+t=[]
+phi=[]
+it = ts.iterations
+
+for i in it:
+    phi_i = ts.get_field('phi',iteration=i,plot=False)
+    # Find the minimum value among all grids for this iteration
+    phi_min = np.min(phi_i[0])
+    phi.append(phi_min)
+    t.append(dt*i)
+
+
+def func(x, v0, tau):
+
+    return v0 * (1-np.exp(-np.array(x) / tau))
+
+
+
+popt, pcov = curve_fit(func, t, phi)
+
+plt.plot(t,phi, label='modelisation')
+plt.plot(t, func(t, *popt), 'r-',label='fit: v0=%5.3f, tau=%5.9f' % (popt[0], popt[1]))
+plt.legend()
+plt.savefig('min_phi_analysis.png')
+
+print('fit parameters between the min(phi) curve over the time and the function v0(1-exp(-t/tau)):')
+print('v0=%5.3f, tau=%5.9f' % (popt[0], popt[1]))
+
+
+tolerance_v0=0.01
+tolerance_tau=0.01
+print("tolerance for v0 = "+ str(tolerance_v0 *100) + '%')
+print("tolerance for tau = "+ str(tolerance_tau*100) + '%')
+
+mean_v0=-151.347
+mean_tau=0.000004351
+
+diff_v0=np.abs((popt[0]-mean_v0)/mean_v0)
+diff_tau=np.abs((popt[1]-mean_tau)/mean_tau)
+
+print("percentage error for v0 = "+ str(diff_v0 *100) + '%')
+print("percentage error for tau = "+ str(diff_tau*100) + '%')
+
+assert (diff_v0 < tolerance_v0) and (diff_tau < tolerance_tau), 'Test spacecraft_charging did not pass'