BLAST-ImpactX · ax3l · Feb 5, 2025 · Feb 5, 2025 · Feb 5, 2025 · Feb 5, 2025
diff --git a/README.md b/README.md
@@ -1,154 +1,131 @@
-# ImpactX
+# Differentiable ImpactX
 
-[![CI Status](https://github.com/BLAST-ImpactX/impactx/actions/workflows/ubuntu.yml/badge.svg)](https://github.com/BLAST-ImpactX/impactx/actions/workflows/ubuntu.yml)
-[![Documentation Status](https://readthedocs.org/projects/impactx/badge/?version=latest)](https://impactx.readthedocs.io)
-[![License ImpactX](https://img.shields.io/badge/license-BSD--3--Clause--LBNL-blue.svg)](https://spdx.org/licenses/BSD-3-Clause-LBNL.html)
-[![Supported Platforms](https://img.shields.io/badge/platforms-linux%20|%20osx%20|%20win-blue)](https://impactx.readthedocs.io/en/latest/install/users.html)  
-[![DOI (source)](https://img.shields.io/badge/DOI%20(source)-10.5281/zenodo.6954922-blue.svg)](https://doi.org/10.5281/zenodo.6954922)
-[![DOI (paper)](https://img.shields.io/badge/DOI%20(paper)-10.18429%2FJACoW--NAPAC2022--TUYE2-blue.svg)](https://doi.org/10.18429/JACoW-NAPAC2022-TUYE2)  
-[![Language: C++17](https://img.shields.io/badge/language-C%2B%2B17-orange.svg)](https://isocpp.org/)
-[![Language: Python](https://img.shields.io/badge/language-Python-orange.svg)](https://python.org/)
+## Prepare Developer Environment (Ubuntu 24.02 w/ Clang/LLVM 16)
 
-ImpactX: an s-based beam dynamics code including space charge effects.
-This is the next generation of the [IMPACT-Z](https://github.com/impact-lbl/IMPACT-Z) code.
-
-## Documentation
-
-In order to learn how to install and run the code, please see the online documentation:
-https://impactx.readthedocs.io
-
-* ImpactX Doxygen: https://impactx.readthedocs.io/en/latest/_static/doxyhtml
-* AMReX Doxygen: https://amrex-codes.github.io/amrex/doxygen
-* WarpX Doxygen: https://warpx.readthedocs.io/en/latest/_static/doxyhtml
-
-## Contributing
+### Once
+```bash
+# Ref: https://github.com/EnzymeAD/enzyme-dev-docker/blob/main/Dockerfile
+sudo apt install llvm-dev clang-16 lld-16 zlib1g-dev libzstd-dev git automake autoconf cmake make lldb-16 ninja-build build-essential libtool llvm-16-dev libclang-16-dev libomp-16-dev libblas-dev libeigen3-dev libboost-dev python3
+python3 -m pip install --break-system-package pipx pathlib2
+python3 -m pipx install lit cmake
+
+export CC="clang-16"
+export CXX="clang++-16"
+
+# Ref: https://enzyme.mit.edu/Installation/
+git clone https://github.com/EnzymeAD/Enzyme ~/src/Enzyme
+cd ~/src/Enzyme
+cmake --fresh -G Ninja -S enzyme -B build -DLLVM_DIR=/usr/lib/llvm-16/lib/cmake/llvm -DLLVM_EXTERNAL_LIT=$(which lit)
+cmake --build build -j 6
+
+cd build
+ninja check-enzyme
+
+# linker and lld and lto enforcement in CMake is hard...
+cd ..
+mkdir mylld
+cd mylld
+ln -s /usr/lib/llvm-16/bin/lld-link lld
+ln -s $(which ld.lld-16) ld.lld  # note: not sufficient yet... somehow hard-coded in compiler detection... use docker
+                                 # manually linking /usr/bin/ld.lld-16 as /usr/bin/ld.lld works...
+export PATH=$PWD:$PATH
+
+# optional: create a venv for Python
+rm -rf ~/src/venv-impactx-enzyme
+python3 -m venv ~/src/venv-impactx-enzyme
+source ~/src/venv-impactx-enzyme/bin/activate
+python3 -m pip install --upgrade pip
+python3 -m pip install --upgrade scipy numpy packaging setuptools[core] wheel pytest pytest-benchmark matplotlib PyQt6
+```
 
-[![AMReX](https://img.shields.io/static/v1?label="runs%20on"&message="AMReX"&color="blueviolet")](https://amrex-codes.github.io/)
 
-Our workflow is described in [CONTRIBUTING.rst](CONTRIBUTING.rst).
+## Compile Options
 
-## Developer Environment
+From the [homepage](https://enzyme.mit.edu/getting_started/UsingEnzyme/):
+> Enzyme supports differentiating C/C++ code through ClangEnzyme and LLDEnzyme as compiler and linker plugins, respectively.
+> Clang gives our plugin more flexibility in adding and ordering optimization passes than LLD and therefore using ClangEnzyme could result in better performance than LLDEnzyme.
+> However, ClangEnzyme can only differentiate one compilation unit at a time and will therefore fail if the function which you try to differentiate calls functions in other compilation units (generally other .c or .cpp files).
+> In these cases we recommend the use of LLDEnzyme in combination with LTO.
 
-Please see our [developer installation section](https://impactx.readthedocs.io/en/latest/install/dependencies.html#install-dependencies) of the documentation for an easy install of our software dependencies.
 
-## Get the Source Code
+## Compile (ClangEnzyme, one TU)
 
-Before you start, you will need a copy of the ImpactX source code:
+### Always
 
 ```bash
-git clone [email protected]:BLAST-ImpactX/impactx.git
-cd impactx
+source ~/src/venv-impactx-enzyme/bin/activate
+alias cmake=$HOME/.local/pipx/venvs/cmake/bin/cmake
+export CC="clang-16"
+export CXX="clang++-16"
+
+# one TU: Clang Plugin
+#   Extra Enzyme options, e.g., print https://enzyme.mit.edu/getting_started/UsingEnzyme/#semantic-options
+#   optional add for verbose output: -mllvm -enzyme-print
+export CXXFLAGS="-fplugin=$HOME/src/Enzyme/build/Enzyme/ClangEnzyme-16.so"
 ```
 
-## Compile
-
+With the active developer env above, inside the ImpactX source dir:
 ```bash
-# find dependencies & configure
-cmake -S . -B build
-
-# compile
-cmake --build build -j 4
+cmake --fresh \
+  -S .        \
+  -B build    \
+  -DImpactX_UNITY_BUILD=ON  \
+  -DImpactX_MPI=OFF       \
+  -DImpactX_COMPUTE=NOACC \
+  -DImpactX_OPENPMD=OFF   \
+  -DCMAKE_LINKER_TYPE=LLD \
+  -DCMAKE_LINKER=/usr/lib/llvm-16/bin/lld-link
+# optional:
+#   -DImpactX_PYTHON=ON -DpyAMReX_IPO=OFF -DImpactX_PYTHON_IPO=OFF
+
+cmake --build build -j 6
+
+# optional:
+cmake --build build -j 6 --target pip_install
 ```
 
-That's all!
-ImpactX binaries are now in `build/bin/`.
-Most people execute these binaries directly or copy them out.
 
-You can inspect and modify build options after running `cmake -S . -B` build with either
+## Compile (LLDEnzyme, multiple TUs)
 
-```bash
-ccmake build
-```
+TODO: redo this part using https://github.com/EnzymeAD/enzyme-dev-docker because ld, lld, ld.ldd with the non-system default is too tricky to get right in CMake (i.e. compiler detection).
 
-or by adding arguments with `-D<OPTION>=<VALUE>` to the first CMake call, e.g.:
+### Always
 
 ```bash
-cmake -S . -B build -DImpactX_COMPUTE=CUDA -DImpactX_MPI=OFF
+export PATH=$HOME/.local/pipx/venvs/cmake/bin:$PATH
+export CC="clang-16"
+export CXX="clang++-16"
+
+# many TU: LDD Plugin
+#   https://github.com/EnzymeAD/Enzyme/blob/main/enzyme/Enzyme/CMakeLists.txt
+#   https://enzyme.mit.edu/getting_started/UsingEnzyme/#semantic-options
+export CXXFLAGS="-fuse-ld=/usr/lib/llvm-16/bin/lld-link -flto"  # -mllvm -enzyme-...
+export LDFLAGS="-fuse-ld=/usr/lib/llvm-16/bin/lld-link -flto -Wl,-mllvm -Wl,-load=$HOME/src/Enzyme/build/Enzyme/LLDEnzyme-16.so -Wl,--load-pass-plugin=$HOME/src/Enzyme/build/Enzyme/LLDEnzyme-16.so"
 ```
 
-### Python Compile
-
+With the active developer env above, inside the ImpactX source dir:
 ```bash
-# find dependencies & configure
-cmake -S . -B build -DImpactX_PYTHON=ON
-
-# compile & install
-cmake --build build -j 4 --target pip_install
+cmake --fresh \
+  -S .        \
+  -B build    \
+  -DImpactX_UNITY_BUILD=ON  \
+  -DImpactX_MPI=OFF       \
+  -DImpactX_COMPUTE=NOACC \
+  -DImpactX_OPENPMD=OFF   \
+  -DCMAKE_LINKER_TYPE=LLD \
+  -DCMAKE_LINKER=/usr/lib/llvm-16/bin/lld-link
+
+# note: -DCMAKE_LINKER_TYPE=LLD appends general lld not necessarily the right version
+#        bend ldd to ldd-16 and /usr/bin/ld.ldd to ld.ldd-16
+# -DCMAKE_INTERPROCEDURAL_OPTIMIZATION=ON
+
+cmake --build build -j 6
 ```
 
-## Run
-
-An executable ImpactX binary with the current compile-time options encoded in its file name will be created in `build/bin/`.
-
-Additionally, a symbolic link named `impactx` can be found in that directory, which points to the last built ImpactX executable.
-
-The command-line syntax for this executable is:
-```console
-Usage: impactx <inputs-file> [some.overwritten.option=value]...
-
-Mandatory arguments (remove the <>):
-  inputs-file     the path to an input file; can be relative to the current
-                  working directory or absolute.
-                  Example: input_fodo.in
-
-Optional arguments (remove the []):
-  options         this can overwrite any line in an inputs-file
-                  Example: quad1.ds=0.5 sbend1.rc=1.5
 
-Examples:
-  In the current working directory, there is a file "input_fodo.in" and the
-  "impactx" executable.
-  The line to execute would look like this:
-    ./impactx input_fodo.in
+## Test/Run
 
-  In the current working directory, there is a file "input_fodo.in" and the
-  executable "impactx" is in a directory that is listed in the "PATH"
-  environment variable.
-  The line to execute would look like this:
-    impactx input_fodo.in
-
-  In the current working directory, there is a file "input_fodo.in" and the
-  "impactx" executable. We want to overwrite the segment length of the beamline
-  element "quad1" that is already defined in it. We also want to change the
-  radius of curvature of the bending magnet "sbend1" to a different value than
-  in the file "input_fodo.in".
-  The line to execute would look like this:
-    ./impactx input_fodo.in quad1.ds=0.5 sbend1.rc=1.5
-```
-
-## Test
-
-In order to run our tests, you need to have a few Python packages installed:
-```console
-python3 -m pip install --upgrade pip
-python3 -m pip install --upgrade build packaging setuptools[core] wheel pytest pytest-benchmark
-python3 -m pip install --upgrade -r tests/python/requirements.txt
-```
-
-You can run all our tests with:
-
-```console
-ctest --test-dir build --output-on-failure
+With the active developer env above:
+```bash
+./build/bin/impactx examples/fodo/input_fodo.in
 ```
-
-Further options:
-* help: `ctest --test-dir build --help`
-* list all tests: `ctest --test-dir build -N`
-* only run tests that have "FODO" in their name: `ctest --test-dir build -R FODO`
-
-## Acknowledgements
-
-This work was supported by the Laboratory Directed Research and Development Program of Lawrence Berkeley National Laboratory under U.S. Department of Energy Contract No. DE-AC02-05CH11231.
-
-ImpactX is supported by the CAMPA collaboration, a project of the U.S. Department of Energy, Office of Science, Office of Advanced Scientific Computing Research and Office of High Energy Physics, Scientific Discovery through Advanced Computing (SciDAC) program.
-
-## Copyright Notice
-
-ImpactX Copyright (c) 2022, The Regents of the University of California, through Lawrence Berkeley National Laboratory (subject to receipt of any required approvals from the U.S. Dept. of Energy).
-All rights reserved.
-
-If you have questions about your rights to use or distribute this software, please contact Berkeley Lab's Intellectual Property Office at [email protected].
-
-Please see the full license agreement in [LICENSE.txt](LICENSE.txt).  
-Please see the notices in [NOTICE.txt](NOTICE.txt).  
-The SPDX license identifier is `BSD-3-Clause-LBNL`.
diff --git a/check_k2_zero.py b/check_k2_zero.py
@@ -0,0 +1,40 @@
+#!/usr/bin/env python3
+#
+# Copyright 2025 ImpactX contributors
+# Authors: Axel Huebl, Chad Mitchell
+# License: BSD-3-Clause-LBNL
+#
+# -*- coding: utf-8 -*-
+
+
+from impactx import my_run
+
+verbose = False
+# mode = "forward"
+mode = "backward"
+inputs_file_beam = "examples/fodo_space_charge/input_fodo_envelope_sc.in"
+
+
+eps = 1e-12
+
+mode = "gradient-free"
+print(f"{mode}:")
+values = my_run(-85.0, -eps, mode, inputs_file_beam, verbose=verbose)
+error_zero_lo = values["error"]
+values = my_run(-85.0, eps, mode, inputs_file_beam, verbose=verbose)
+error_zero_hi = values["error"]
+print(f"  error_zero_lo={error_zero_lo}, error_zero_hi={error_zero_hi}")
+print(f"  FD gradient: derror_dq2_k={(error_zero_hi - error_zero_lo) / (2 * eps)}")
+
+values = my_run(-85.0, 0.0, mode, inputs_file_beam, verbose=verbose)
+print(f"  error_forward={values['error']}")
+
+mode = "forward"
+print(f"\n{mode}:")
+values = my_run(-85.0, 0.0, mode, inputs_file_beam, verbose=verbose)
+print(f"  error_forward={values['error']} derror_dq2_k={values['derror_dq2_k']}")
+
+mode = "backward"
+print(f"\n{mode}:")
+values = my_run(-85.0, 0.0, mode, inputs_file_beam, verbose=verbose)
+print(f"  error_forward={values['error']} derror_dq2_k={values['derror_dq2_k']}")
diff --git a/cmake/dependencies/ABLASTR.cmake b/cmake/dependencies/ABLASTR.cmake
@@ -61,6 +61,8 @@ macro(find_ablastr)
         set(ABLASTR_FFT ${ImpactX_FFT} CACHE BOOL "" FORCE)
         set(AMReX_FFT ${ImpactX_FFT} CACHE BOOL "" FORCE)
 
+        set(AMReX_TINY_PROFILE OFF CACHE BOOL "")
+
         set(WarpX_APP OFF CACHE BOOL "" FORCE)
         set(WarpX_LIB OFF CACHE BOOL "" FORCE)
         set(WarpX_QED OFF CACHE BOOL "" FORCE)

diff --git a/examples/fodo_space_charge/input_fodo_envelope_sc.in b/examples/fodo_space_charge/input_fodo_envelope_sc.in
@@ -1,3 +1,5 @@
+impactx.verbose = 0
+
 ###############################################################################
 # Particle Beam(s)
 ###############################################################################
@@ -15,37 +17,9 @@ 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