ECP-WarpX · cemitch99 · Jun 7, 2022 · Dec 5, 2022 · Sep 8, 2023 · Sep 8, 2023
diff --git a/docs/source/usage/parameters.rst b/docs/source/usage/parameters.rst
@@ -625,6 +625,29 @@ Numerics and algorithms
     This is in-development.
     At the moment, this flag only activates coordinate transformations and charge deposition.
 
+* ``algo.mlmg_relative_tolerance`` (``float``, optional, default: ``1.e-7``)
+    The relative precision with which the electrostatic space-charge fields should be calculated.
+    More specifically, the space-charge fields are computed with an iterative Multi-Level Multi-Grid (MLMG) solver.
+    This solver can fail to reach the default precision within a reasonable time.
+
+* ``algo.mlmg_absolute_tolerance`` (``float``, optional, default: ``0``, which means: ignored)
+    The absolute tolerance with which the space-charge fields should be calculated in units of V/m^2.
+    More specifically, the acceptable residual with which the solution can be considered converged.
+    In general this should be left as the default, but in cases where the simulation state changes very
+    little between steps it can occur that the initial guess for the MLMG solver is so close to the
+    converged value that it fails to improve that solution sufficiently to reach the
+    mlmg_relative_tolerance value."
+
+* ``algo.mlmg_max_iters`` (``integer``, optional, default: ``100``)
+    Maximum number of iterations used for MLMG solver for space-charge fields calculation.
+    In case if MLMG converges but fails to reach the desired self_fields_required_precision,
+    this parameter may be increased.
+
+* ``algo.mlmg_verbosity`` (``integer``, optional, default: ``1``)
+    The verbosity used for MLMG solver for space-charge fields calculation.
+    Currently MLMG solver looks for verbosity levels from 0-5.
+    A higher number results in more verbose output.
+
 .. _running-cpp-parameters-diagnostics:
 
 Diagnostics and output

diff --git a/docs/source/usage/python.rst b/docs/source/usage/python.rst
@@ -66,6 +66,39 @@ General
       This is in-development.
       At the moment, this flag only activates coordinate transformations and charge deposition.
 
+   .. py:property:: mlmg_relative_tolerance
+
+      Default: ``1.e-7``
+
+      The relative precision with which the electrostatic space-charge fields should be calculated.
+      More specifically, the space-charge fields are computed with an iterative Multi-Level Multi-Grid (MLMG) solver.
+      This solver can fail to reach the default precision within a reasonable time.
+
+   .. py:property:: mlmg_absolute_tolerance
+
+      Default: ``0``, which means: ignored
+
+      The absolute tolerance with which the space-charge fields should be calculated in units of :math:`V/m^2`.
+      More specifically, the acceptable residual with which the solution can be considered converged.
+      In general this should be left as the default, but in cases where the simulation state changes very
+      little between steps it can occur that the initial guess for the MLMG solver is so close to the
+      converged value that it fails to improve that solution sufficiently to reach the ``mlmg_relative_tolerance`` value.
+
+   .. py:property:: mlmg_max_iters
+
+      Default: ``100``
+
+      Maximum number of iterations used for MLMG solver for space-charge fields calculation.
+      In case if MLMG converges but fails to reach the desired self_fields_required_precision, this parameter may be increased.
+
+   .. py:property:: mlmg_verbosity
+
+      Default: ``1``
+
+      The verbosity used for MLMG solver for space-charge fields calculation.
+      Currently MLMG solver looks for verbosity levels from 0-5.
+      A higher number results in more verbose output.
+
    .. py:property:: diagnostics
 
       Enable (``True``) or disable (``False``) diagnostics generally (default: ``True``).
@@ -311,7 +344,7 @@ Particles
 
    .. py:method:: set_energy_MeV(energy_MeV)
 
-      Write-only: Set reference particle energy.
+      Write-only: Set reference particle kinetic energy.
 
    .. py:method:: load_file(madx_file)
 

diff --git a/examples/ipac2018_mayes/Ex_Mayes.dat b/examples/ipac2018_mayes/Ex_Mayes.dat
diff --git a/examples/ipac2018_mayes/Ez_Mayes.dat b/examples/ipac2018_mayes/Ez_Mayes.dat
diff --git a/examples/ipac2018_mayes/Mayes_Fig1a_Script.sh b/examples/ipac2018_mayes/Mayes_Fig1a_Script.sh
@@ -0,0 +1,17 @@
+set xtics font  'helvetica,25'
+set ytics font 'helvetica,25'
+set xlabel 'x/\sigmax' font 'helvetica,30' offset 0,-1
+set ylabel 'Ex (MV/m)' font 'helvetica,30' offset -1,0
+set lmargin 12
+set bmargin 6
+set xrange [-6:6]
+set yrange [-10:10]
+set nokey
+set mxtics 5
+set mytics 5
+f(x)=8.0*exp(-x**2/2.0)
+plot f(x) w filledcu y1=0.0 lt 3 fs transparent solid 0.25
+replot 'Ex_Mayes.dat' u 1:2 w l lw 2 lt 1
+replot 'Ex_Mayes.dat' u 1:3 w l lw 2 lt 6
+replot 'Ex_Mayes.dat' u 1:4 w l lw 2 lt 8
+replot 'Ex_Mayes.dat' u 1:5 w l lw 2 lt 7
diff --git a/examples/ipac2018_mayes/Mayes_Fig1b_Script.sh b/examples/ipac2018_mayes/Mayes_Fig1b_Script.sh
@@ -0,0 +1,17 @@
+set xtics font  'helvetica,25'
+set ytics font 'helvetica,25'
+set xlabel 'z/\sigmaz' font 'helvetica,30' offset 0,-1
+set ylabel 'Ez (MV/m)' font 'helvetica,30' offset -1,0
+set lmargin 12
+set bmargin 6
+set xrange [-6:6]
+set yrange [-10:10]
+set nokey
+set mxtics 5
+set mytics 5
+f(x)=8.0*exp(-x**2/2.0)
+plot f(x) w filledcu y1=0.0 lt 3 fs transparent solid 0.25
+replot 'Ez_Mayes.dat' u 1:2 w l lw 2 lt 1
+replot 'Ez_Mayes.dat' u 1:3 w l lw 2 lt 6
+replot 'Ez_Mayes.dat' u 1:4 w l lw 2 lt 8
+replot 'Ez_Mayes.dat' u 1:5 w l lw 2 lt 7
diff --git a/examples/ipac2018_mayes/input_expanding_ipac2018.in b/examples/ipac2018_mayes/input_expanding_ipac2018.in
@@ -0,0 +1,38 @@
+###############################################################################
+# Particle Beam(s)
+###############################################################################
+beam.npart = 10000  # outside tests, use 1e5 or more
+beam.units = static
+beam.energy = 9.48900105  # KE corresponding to 10 MeV total energy
+beam.charge = 1.0e-9  # 1 nC
+beam.particle = electron
+beam.distribution = gaussian
+beam.sigmaX = 1.0e-3  # 1 mm
+beam.sigmaY = 1.0e-3  # 1 mm
+beam.sigmaT = 1.00130816210e-4  # corresponding to sig_z = 0.1 mm
+beam.sigmaPx = 0.0
+beam.sigmaPy = 0.0
+beam.sigmaPt = 0.0
+
+
+###############################################################################
+# Beamline: lattice elements and segments
+###############################################################################
+lattice.elements = monitor drift1 monitor
+lattice.nslice = 40
+
+drift1.type = drift
+drift1.ds = 1.0  # 1 m drift
+
+monitor.type = beam_monitor
+monitor.backend = h5
+
+
+###############################################################################
+# Algorithms
+###############################################################################
+algo.particle_shape = 2
+algo.space_charge = true
+
+amr.n_cell = 56 56 48
+geometry.prob_relative = 3.0
diff --git a/examples/ipac2018_mayes/plot_expanding_ipac2018.py b/examples/ipac2018_mayes/plot_expanding_ipac2018.py
@@ -0,0 +1,56 @@
+#!/usr/bin/env python3
+#
+# Copyright 2022-2023 ImpactX contributors
+# Authors: Axel Huebl, Chad Mitchell
+# License: BSD-3-Clause-LBNL
+#
+
+
+import matplotlib.pyplot as plt
+from matplotlib.ticker import MaxNLocator
+import openpmd_api as io
+import pandas as pd
+
+# collect final beam
+series = io.Series("diags/openPMD/monitor.h5", io.Access.read_only)
+last_step = list(series.iterations)[-1]
+final_beam = series.iterations[last_step].particles["beam"].to_df()
+
+# scaling to figure units
+beta = 0.9986935469557160
+gamma = 19.569511835591836
+ErMeV = 0.510998950
+scalex = 1.0e3  # x [m] -> x [mm]
+scalet = -beta * 1.0e3  # ct [m] -> z [m]
+scalepx = beta * gamma * ErMeV  # px/p0 -> px [MeV/c]
+scalept = -gamma * ErMeV  # pt/p0 -> pz [MeV/c]
+
+# beam phase space scatter plots
+num_plots_per_row = 2
+fig, axs = plt.subplots(1, num_plots_per_row, figsize=(10, 4.0))
+
+# plot final x-px distribution
+ax = axs[(0)]
+ax.scatter(
+    final_beam.position_x.multiply(scalex),
+    final_beam.momentum_x.multiply(scalepx),
+    s=0.5,
+)
+ax.tick_params(labelsize=12)
+ax.set_xlabel(r"$x$ [mm]", fontsize=18)
+ax.set_ylabel(r"$p_x$ [MeV/c]", fontsize=18)
+
+# plot final z-pz distribution
+ax = axs[(1)]
+ax.scatter(
+    final_beam.position_t.multiply(scalet),
+    final_beam.momentum_t.multiply(scalept),
+    s=0.5,
+)
+ax.tick_params(labelsize=12)
+ax.set_xlabel(r"$z$ [mm]", fontsize=18)
+ax.set_ylabel(r"$p_z$ [MeV/c]", fontsize=18)
+
+# store plots
+fig.tight_layout()
+plt.savefig("expanding_scatter.png")
diff --git a/examples/ipac2018_mayes/run_expanding_ipac2018.py b/examples/ipac2018_mayes/run_expanding_ipac2018.py
@@ -0,0 +1,110 @@
+#!/usr/bin/env python3
+#
+# Copyright 2022 ImpactX contributors
+# Authors: Axel Huebl, Christopher E. Mayes
+# License: BSD-3-Clause-LBNL
+#
+# -*- coding: utf-8 -*-
+
+import numpy as np
+
+import amrex.space3d as amr
+from impactx import ImpactX, RefPart, distribution, elements
+
+sim = ImpactX()
+
+# set numerical parameters and IO control
+sim.n_cell = [56, 56, 48]
+sim.particle_shape = 2  # B-spline order
+sim.space_charge = True
+sim.dynamic_size = True
+sim.prob_relative = 3.0
+
+# beam diagnostics
+# sim.diagnostics = False  # benchmarking
+sim.slice_step_diagnostics = False
+
+# domain decomposition & space charge mesh
+sim.init_grids()
+
+# load a cold 10 MeV electron beam
+energy_MeV = 10.0  # reference energy (total)
+mass_MeV = 0.510998950  # electron mass in MeV/c^2
+bunch_charge_C = 1.0e-9  # charge in C
+npart = int(10000)  # number of macro particles
+
+#   reference particle
+ref = sim.particle_container().ref_particle()
+ref.set_charge_qe(-1.0).set_mass_MeV(mass_MeV).set_energy_MeV(energy_MeV)
+
+#   particle bunch
+r = 0.1  # aspect ratio = sigma_z / sigma_perp:  range 0.01 to 10
+sigma_r = 1.0e-3  # fixed at 1 mm
+sigma_z = r * sigma_r
+gamma = energy_MeV / mass_MeV
+beta = (1.0 - (1.0 / gamma) ** 2) ** 0.5
+c0 = 2.99792458e8  # speed of light in m/s
+sigma_t = sigma_z / beta  # recall t is implicitly scaled by c0
+print(f"sigma_t={sigma_t}m")
+
+distr = distribution.Gaussian(
+    sigmaX=sigma_r,
+    sigmaY=sigma_r,
+    sigmaT=sigma_t,
+    sigmaPx=0.0,
+    sigmaPy=0.0,
+    sigmaPt=0.0,
+)
+sim.add_particles(bunch_charge_C, distr, npart)
+
+# add beam diagnostics
+monitor = elements.BeamMonitor("monitor", backend="h5")
+
+# design the accelerator lattice
+sim.lattice.extend([monitor, elements.Drift(ds=1.0, nslice=30), monitor])
+
+# run simulation
+sim.evolve()
+
+# calculate phase space
+#   hack:
+import matplotlib.pyplot as plt
+
+num_plots_per_row = 2
+f, axs = plt.subplots(1, num_plots_per_row, figsize=(7, 2))
+ax_x_px = axs[0]
+ax_z_pz = axs[1]
+
+pc = sim.particle_container()
+lev = pc.GetParticles(0)
+for tile_ind, pt in lev.items():
+    # positions + id + cpuid
+    aos = pt.GetArrayOfStructs()
+    aos_arr = np.array(aos, copy=False)
+
+    # momentum & particle weight
+    real_arrays = pt.GetStructOfArrays().GetRealData()
+    px = np.array(real_arrays[0], copy=False)
+    pz = np.array(real_arrays[2], copy=False)
+
+    print(f"tile_ind={tile_ind}, pt={pt}")
+    print(f"aos_arr={aos_arr}, aos_arr.shape={aos_arr.shape}")
+    print(f"px={px}, px.shape={px.shape}")
+
+    ax_x_px.scatter(aos_arr[()]["x"], px)
+    ax_z_pz.scatter(aos_arr[()]["z"], pz)
+plt.show()
+
+# MPI reduce phase space
+#   TODO SUM up histograms via MPI
+
+# plot phase space
+# import matplotlib.pyplot as plt
+# f = plt.figure()
+# ax = plt.gca()
+# ax.scatter()
+
+# clean shutdown
+#   note: timers turned stop here
+del sim
+amr.finalize()