Space charge 2D envelope push

docs/source/usage/examples.rst

@@ -52,6 +52,7 @@ Space Charge
    examples/cfchannel/README.rst
    examples/kurth/README.rst
    examples/epac2004_benchmarks/README.rst
+   examples/fodo_space_charge/README.rst
 
 Coherent Synchrotron Radiation (CSR)
 """"""""""""""""""""""""""""""""""""

docs/source/usage/how_to_run.rst

@@ -13,13 +13,13 @@ ImpactX can be run using any of three distinct tracking modes.  ImpactX's most p
 Additionally, ImpactX provides two simplified tracking modes to aid scientists through every step, from beamline inception to operation:
 tracking of the beam envelope (6x6 covariance matrix) through linearized transport maps, or only tracking of the reference particle orbit.
 
-================== =============== =============== ==================
+================== =============== =============== ===================
 Mode               Use Case        Generality      Collective Effects
-================== =============== =============== ==================
-Particle Tracking  Full Dynamics   Most general    Supported
-Envelope Tracking  Rapid Scans     Linearized      `Soon <https://github.com/ECP-WarpX/impactx/issues/826>`__
+================== =============== =============== ===================
+Particle Tracking  Full Dynamics   Most general    Supported (3D only)
+Envelope Tracking  Rapid Scans     Linearized      Supported (2D only)
 Reference Tracking Early Design    Reference orbit No
-================== =============== =============== ==================
+================== =============== =============== ===================
 
 
 .. _usage_run-user-interface:

docs/source/usage/parameters.rst

@@ -42,6 +42,9 @@ Initial Beam Distributions
 * ``beam.charge`` (``float``, in C)
   bunch charge
 
+* ``beam.current`` (``float``, in A)
+  beam current, used only if ``algo.space_charge = "2D"``
+
 * ``beam.particle`` (``string``)
   particle type: currently either ``electron``, ``positron`` or ``proton``
 
@@ -714,12 +717,24 @@ Space Charge
 Space charge kicks are applied in between slices of thick :ref:`lattice elements <running-cpp-parameters-lattice>`.
 See there ``nslice`` option on lattice elements for slicing.
 
-* ``algo.space_charge`` (``boolean``, optional, default: ``false``)
+* ``algo.space_charge`` (``string``, optional)
+
+  The physical model of space charge used.
+
+  ImpactX uses an AMReX grid of boxes to organize and parallelize space charge simulation domain.
+  These boxes also contain a field mesh, if space charge calculations are enabled.
+
+  Options:
+
+  * ``"false"`` (default): space charge effects are not calculated.
+
+  * ``"2D"``: Space charge forces are computed in the plane ``(x,y)`` transverse to the reference particle velocity, assuming the beam is long and unbunched.
+
+    Currently, this model is supported only in envelope mode (when ``algo.track = "envelope"``).
 
-  Whether to calculate space charge effects.
+  * ``"3D"``: Space charge forces are computed in three dimensions, assuming the beam is bunched.
 
-ImpactX uses an AMReX grid of boxes to organize and parallelize space charge simulation domain.
-These boxes also contain a field mesh, if space charge calculations are enabled.
+    Currently, this model is supported only in particle mode (when ``algo.track = "particles"``).
 
 * ``amr.n_cell`` (3 integers) optional (default: 1 `blocking_factor <https://amrex-codes.github.io/amrex/docs_html/GridCreation.html>`__ per MPI process)
 

docs/source/usage/python.rst

@@ -62,9 +62,19 @@ Collective Effects & Overall Simulation Parameters
 
    .. py:property:: space_charge
 
-      Enable (``True``) or disable (``False``) space charge calculations (default: ``False``).
+      The physical model of space charge used.
+
+      Options:
+
+      * ``False`` (default): space charge effects are not calculated.
+
+      * ``"2D"``: Space charge forces are computed in the plane ``(x,y)`` transverse to the reference particle velocity, assuming the beam is long and unbunched.
+
+        Currently, this model is supported only in envelope mode (when ``algo.track = "envelope"``).
+
+      * ``"3D"``: Space charge forces are computed in three dimensions, assuming the beam is bunched.
 
-      Whether to calculate space charge effects.
+        Currently, this model is supported only in particle mode (when ``algo.track = "particles"``).
 
    .. py:property:: poisson_solver
 
@@ -180,13 +190,14 @@ Collective Effects & Overall Simulation Parameters
       :param distr: distribution function to draw from (object from :py:mod:`impactx.distribution`)
       :param int npart: number of particles to draw
 
-   .. py:method:: init_envelope(ref, distr)
+   .. py:method:: init_envelope(ref, distr, intensity=None)
 
       Envelope tracking mode:
       Create a 6x6 covariance matrix from a distribution and then initialize the the simulation for envelope tracking relative to a reference particle.
 
       :param ref: the reference particle (object from :py:class:`impactx.RefPart`)
       :param distr: distribution function (object from :py:mod:`impactx.distribution`)
+      :param float intensity: the beam intensity, given as bunch charge (C) for 3D or beam current (A) for 2D space charge
 
    .. py:method:: particle_container()
 

examples/CMakeLists.txt

@@ -106,7 +106,7 @@ function(add_impactx_test name input is_mpi analysis_script plot_script)
                         amrex.throw_exception = 1
                         amrex.signal_handling = 0
                         impactx.always_warn_immediately=1
-                        impactx.abort_on_warning_threshold=low
+                        #impactx.abort_on_warning_threshold=low
                         ${INPUTS_ARGS}
                  WORKING_DIRECTORY ${THIS_WORKING_DIR}
         )
@@ -1269,3 +1269,19 @@ add_impactx_test(examples-scraping.py
     examples/scraping_beam/analysis_scraping.py
     OFF  # no plot script yet
 )
+
+# FODO cell with 2D space charge using envelope tracking ######################
+#
+# with space charge
+add_impactx_test(FODO.envelope.sc
+    examples/fodo_space_charge/input_fodo_envelope_sc.in
+    ON   # ImpactX MPI-parallel
+    examples/fodo_space_charge/analysis_fodo_envelope_sc.py
+    OFF  # no plot script yet
+)
+add_impactx_test(FODO.envelope.sc.py
+    examples/fodo_space_charge/run_fodo_envelope_sc.py
+    OFF   # ImpactX MPI-parallel
+    examples/fodo_space_charge/analysis_fodo_envelope_sc.py
+    OFF  # no plot script yet
+)

examples/fodo_space_charge/README.rst

@@ -0,0 +1,53 @@
+.. _examples-fodo-envelope-sc:
+
+FODO Cell with 2D Space Charge using Envelope Tracking
+======================================================
+
+This example illustrates a 0.5 A proton beam with a kinetic energy of 6.7 MeV in a FODO cell,
+with 2D space charge included.  The parameters are those described in:
+
+R.D. Ryne et al, `"A Test Suite of Space-Charge Problems for Code Benchmarking," <https://accelconf.web.cern.ch/e04/PAPERS/WEPLT047.PDF>`__
+in Proc. EPAC 2004, Lucerne, Switzerland:  KV Beam in a FODO Channel
+
+The purpose of this example is to illustrate the use of envelope tracking mode with 2D space charge.
+
+The second moments of the particle distribution after the FODO cell should coincide with the second moments of the particle distribution before the FODO cell, to within the level expected due to noise due to statistical sampling.
+
+In this test, the initial and final values of :math:`\lambda_x`, :math:`\lambda_y`, :math:`\lambda_t`, :math:`\epsilon_x`, :math:`\epsilon_y`, and :math:`\epsilon_t` must agree with nominal values.
+
+
+Run
+---
+
+This example can be run **either** as:
+
+* **Python** script: ``python3 run_fodo_envelope_sc.py`` or
+* ImpactX **executable** using an input file: ``impactx input_fodo_envelope_sc.in``
+
+For `MPI-parallel <https://www.mpi-forum.org>`__ runs, prefix these lines with ``mpiexec -n 4 ...`` or ``srun -n 4 ...``, depending on the system.
+
+.. tab-set::
+
+   .. tab-item:: Python: Script
+
+       .. literalinclude:: run_fodo_envelope_sc.py
+          :language: python3
+          :caption: You can copy this file from ``examples/fodo_space_charge/run_fodo_envelope_sc.py``.
+
+   .. tab-item:: Executable: Input File
+
+       .. literalinclude:: input_fodo_envelope_sc.in
+          :language: ini
+          :caption: You can copy this file from ``examples/fodo_space_charge/input_fodo_envelope_sc.in``.
+
+
+Analyze
+-------
+
+We run the following script to analyze correctness:
+
+.. dropdown:: Script ``analysis_fodo_envelope_sc.py``
+
+   .. literalinclude:: analysis_fodo_envelope_sc.py
+      :language: python3
+      :caption: You can copy this file from ``examples/fodo_space_charge/analysis_fodo_envelope_sc.py``.

examples/fodo_space_charge/analysis_fodo_envelope_sc.py

@@ -0,0 +1,116 @@
+#!/usr/bin/env python3
+#
+# Copyright 2022-2023 ImpactX contributors
+# Authors: Axel Huebl, Chad Mitchell
+# License: BSD-3-Clause-LBNL
+#
+
+import glob
+import re
+
+import numpy as np
+import pandas as pd
+
+
+def read_file(file_pattern):
+    for filename in glob.glob(file_pattern):
+        df = pd.read_csv(filename, delimiter=r"\s+")
+        if "step" not in df.columns:
+            step = int(re.findall(r"[0-9]+", filename)[0])
+            df["step"] = step
+        yield df
+
+
+def read_time_series(file_pattern):
+    """Read in all CSV files from each MPI rank (and potentially OpenMP
+    thread). Concatenate into one Pandas dataframe.
+
+    Returns
+    -------
+    pandas.DataFrame
+    """
+    return pd.concat(
+        read_file(file_pattern),
+        axis=0,
+        ignore_index=True,
+    )  # .set_index('id')
+
+
+# read reduced diagnostics
+rbc = read_time_series("diags/reduced_beam_characteristics.*")
+
+s = rbc["s"]
+sig_x = rbc["sig_x"]
+sig_y = rbc["sig_y"]
+sig_t = rbc["sig_t"]
+emittance_x = rbc["emittance_x"]
+emittance_y = rbc["emittance_y"]
+emittance_t = rbc["emittance_t"]
+
+sig_xi = sig_x.iloc[0]
+sig_yi = sig_y.iloc[0]
+sig_ti = sig_t.iloc[0]
+emittance_xi = emittance_x.iloc[0]
+emittance_yi = emittance_y.iloc[0]
+emittance_ti = emittance_t.iloc[0]
+
+length = len(s) - 1
+
+sf = s.iloc[length]
+sig_xf = sig_x.iloc[length]
+sig_yf = sig_y.iloc[length]
+sig_tf = sig_t.iloc[length]
+emittance_xf = emittance_x.iloc[length]
+emittance_yf = emittance_y.iloc[length]
+emittance_tf = emittance_t.iloc[length]
+
+
+print("Initial Beam:")
+print(f"  sigx={sig_xi:e} sigy={sig_yi:e} sigt={sig_ti:e}")
+print(
+    f"  emittance_x={emittance_xi:e} emittance_y={emittance_yi:e} emittance_t={emittance_ti:e}"
+)
+
+atol = 0.0  # ignored
+rtol = 1.0e-3  # from random sampling of a smooth distribution
+print(f"  rtol={rtol} (ignored: atol~={atol})")
+
+assert np.allclose(
+    [sig_xi, sig_yi, sig_ti, emittance_xi, emittance_yi, emittance_ti],
+    [
+        8.590000e-04,
+        8.590000e-04,
+        7.071068e-07,
+        1.000000e-06,
+        1.000000e-06,
+        1.000000e-12,
+    ],
+    rtol=rtol,
+    atol=atol,
+)
+
+
+print("")
+print("Final Beam:")
+print(f"  sigx={sig_xf:e} sigy={sig_yf:e} sigt={sig_tf:e}")
+print(
+    f"  emittance_x={emittance_xf:e} emittance_y={emittance_yf:e} emittance_t={emittance_tf:e}"
+)
+
+atol = 0.0  # ignored
+rtol = 1.0e-3  # from random sampling of a smooth distribution
+print(f"  rtol={rtol} (ignored: atol~={atol})")
+
+assert np.allclose(
+    [sig_xf, sig_yf, sig_tf, emittance_xf, emittance_yf, emittance_tf],
+    [
+        8.590000e-04,
+        8.590000e-04,
+        4.140854e-05,
+        1.000000e-06,
+        1.000000e-06,
+        1.000000e-12,
+    ],
+    rtol=rtol,
+    atol=atol,
+)

examples/fodo_space_charge/input_fodo_envelope_sc.in

@@ -0,0 +1,51 @@
+###############################################################################
+# Particle Beam(s)
+###############################################################################
+beam.kin_energy = 6.7
+beam.current = 0.5
+beam.particle = proton
+beam.distribution = kvdist_from_twiss
+beam.alphaX = 2.4685083
+beam.alphaY = -beam.alphaX
+beam.alphaT = 0.0
+beam.betaX = 0.737881
+beam.betaY = 0.737881
+beam.betaT = 0.5
+beam.emittX = 1.0e-6
+beam.emittY = beam.emittX
+beam.emittT = 1.0e-12
+
+###############################################################################
+# Beamline: lattice elements and segments
+###############################################################################
+lattice.elements = monitor drift1 quad1 drift2 quad2 drift1 monitor
+lattice.nslice = 50
+
+monitor.type = beam_monitor
+monitor.backend = h5
+
+drift1.type = drift
+drift1.ds = 7.44e-2
+
+quad1.type = quad
+quad1.ds = 6.10e-2
+quad1.k = -103.12574100336
+
+drift2.type = drift
+drift2.ds = 14.88e-2
+
+quad2.type = quad
+quad2.ds = 6.10e-2
+quad2.k = -quad1.k
+
+###############################################################################
+# Algorithms
+###############################################################################
+algo.particle_shape = 2
+algo.track = "envelope"
+algo.space_charge = 2D
+
+###############################################################################
+# Diagnostics
+###############################################################################
+diag.slice_step_diagnostics = true