diff --git a/README.md b/README.md
index 4445ab533..5c438e774 100644
--- 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 git@github.com: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 IPO@lbl.gov.
-
-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
new file mode 100755
index 000000000..ea8d867c1
--- /dev/null
+++ 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
index 5b4666734..20b6c19d8 100644
--- 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
index 0ee52f44f..a12183718 100644
--- 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
diff --git a/match_beam.py b/match_beam.py
new file mode 100755
index 000000000..ed98e0384
--- /dev/null
+++ b/match_beam.py
@@ -0,0 +1,168 @@
+#!/usr/bin/env python3
+#
+# Copyright 2025 ImpactX contributors
+# Authors: Axel Huebl, Chad Mitchell
+# License: BSD-3-Clause-LBNL
+#
+# -*- coding: utf-8 -*-
+
+import time
+
+import matplotlib.pyplot as plt
+import numpy as np
+from scipy.optimize import minimize
+
+from impactx import my_run
+
+verbose = False
+# mode = "forward"
+mode = "backward"
+inputs_file_beam = "examples/fodo_space_charge/input_fodo_envelope_sc.in"
+
+history = {}
+runtime = {}
+optimizer = None
+grads = False
+
+
+class ProcessTimer:
+    def __init__(self):
+        self.elapsed_time = 0
+
+    def __enter__(self):
+        self.start_time = time.process_time_ns()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        end_time = time.process_time_ns()
+        self.elapsed_time = end_time - self.start_time
+
+
+def objective(parameters: tuple) -> float:
+    """
+    A function that is evaluated by the optimizer.
+
+    Parameters
+    ----------
+    parameters: tuple
+      quadrupole strengths k of quad 1 and quad 2.
+
+    Returns
+    -------
+    The L2 norm of alpha and beta of the beam at the end of the
+    simulation.
+    """
+    global grads
+
+    if verbose:
+        print(f"Run objective with parameters={parameters}...")
+
+    if not grads:
+        use_mode = "gradient-free"
+    else:
+        use_mode = mode
+
+    q1_k, q2_k = parameters
+
+    with ProcessTimer() as time_ns:
+        values = my_run(q1_k, q2_k, use_mode, inputs_file_beam, verbose)
+    # print(f"optimizer (grads={grads}): {time_ns.elapsed_time}ns")
+    error = values["error"]
+    # print(q1_k, q2_k, error, values)
+
+    history[optimizer].append([q1_k, q2_k])
+    runtime[optimizer] = min(time_ns.elapsed_time, runtime[optimizer])
+
+    if verbose:
+        print(f"error={error}, q1_k={q1_k}, q2_k={q2_k}")
+
+    if np.isnan(error):
+        error = 1.0e99
+
+    if grads:
+        return (error, [values["derror_dq1_k"], values["derror_dq2_k"]])
+    else:
+        return error
+
+
+def optimize_and_plot(opti, jac=None, plot=True):
+    global grads, history, optimizer, runtime
+    grads = jac == True  # noqa
+    optimizer = opti
+    history[optimizer] = []
+    runtime[optimizer] = 9e99
+    res = minimize(
+        objective,
+        initial_quad_strengths,
+        method=optimizer,
+        jac=jac,
+        tol=1.0e-8,
+        options=options,
+        bounds=bounds,
+    )
+    ln = len(history[optimizer])
+    print(optimizer, ln)
+    q1_k, q2_k = zip(*history[optimizer])
+    if plot:
+        opt_str = optimizer
+        if optimizer == "CG":
+            opt_str = "Conjugate Gradient"
+        plt.plot(q1_k, q2_k, label=f"{opt_str}: {ln}x")
+    print("  Optimal parameters for k:", res.x)
+    print("  L2 norm of alpha & beta at the optimum:", res.fun)
+    print(f"  Min. runtime: {runtime[optimizer] / ln}ns")
+
+
+if __name__ == "__main__":
+    # Initial guess for the quadrople strengths
+    initial_quad_strengths = np.array([-85, 80])
+
+    # Bounds for values to test: (min, max)
+    positive = (0, None)
+    negative = (None, 0)
+    bounds = [negative, positive]
+
+    # optimizer specific values
+    #   https://docs.scipy.org/doc/scipy/reference/optimize.minimize-neldermead.html
+    #   https://docs.scipy.org/doc/scipy/reference/optimize.minimize-lbfgsb.html
+    options = {
+        "maxiter": 1000,
+    }
+
+    fig = plt.figure(figsize=(4.5, 2.2))
+
+    # Call the optimizer
+    optimize_and_plot("Nelder-Mead", False)
+    optimize_and_plot("CG", "2-point", False)
+
+    # gradient-based: CG, BFGS, Newton-CG, L-BFGS-B, TNC, SLSQP, dogleg, trust-ncg, trust-krylov, trust-exact and trust-constr
+    optimize_and_plot("CG", True)
+
+    # analytical result:
+    #   q1_k = -103.12574100336
+    #   q2_k = -q1_k
+
+    fig.gca().annotate(
+        "initial quadrupole strengths",  # Annotation text
+        xy=(-85, 80),  # Point to annotate (arrow points here)
+        xytext=(-95.8, 62),  # Text position (arrow starts here)
+        arrowprops=dict(
+            facecolor="black", shrink=0.1, width=1.5, headwidth=4, headlength=6
+        ),
+    )
+    fig.gca().annotate(
+        "solution",  # Annotation text
+        xy=(-103.12574100336, 103.12574100336),  # Point to annotate (arrow points here)
+        xytext=(-106, 113),  # Text position (arrow starts here)
+        arrowprops=dict(
+            facecolor="black", shrink=0.1, width=1.5, headwidth=4, headlength=6
+        ),
+    )
+
+    plt.gca().set_xlim(-108, -80)
+    plt.gca().set_ylim(60, 120)
+    plt.xlabel(r"$k_1$  [1/m$^2$]")
+    plt.ylabel(r"$k_2$  [1/m$^2$]")
+    plt.legend(loc="upper right")
+    plt.tight_layout()
+    plt.show()
diff --git a/match_runtimes.txt b/match_runtimes.txt
new file mode 100644
index 000000000..d827f53c5
--- /dev/null
+++ b/match_runtimes.txt
@@ -0,0 +1,41 @@
+BW
+
+Nelder-Mead 167
+  Optimal parameters for k: [-103.0293668   102.82694643]
+  L2 norm of alpha & beta at the optimum: 2.1624999901619736e-05
+  Min. runtime: 6999.347305389221ns
+/home/axel/src/impactx/./match_beam.py:94: RuntimeWarning: Method CG cannot handle bounds.
+  res = minimize(
+CG 150
+  Optimal parameters for k: [-103.02936496  102.82694212]
+  L2 norm of alpha & beta at the optimum: 2.1624999956742072e-05
+  Min. runtime: 7764.473333333333ns
+/home/axel/src/impactx/./match_beam.py:94: RuntimeWarning: Method CG cannot handle bounds.
+  res = minimize(
+CG 37
+  Optimal parameters for k: [-103.02936686  102.82694637]
+  L2 norm of alpha & beta at the optimum: 2.1624999902094705e-05
+  Min. runtime: 34293.21621621621ns
+
+
+
+
+
+FW
+
+Nelder-Mead 167
+  Optimal parameters for k: [-103.0293668   102.82694643]
+  L2 norm of alpha & beta at the optimum: 2.1624999901619736e-05
+  Min. runtime: 6997.185628742515ns
+/home/axel/src/impactx/./match_beam.py:94: RuntimeWarning: Method CG cannot handle bounds.
+  res = minimize(
+CG 150
+  Optimal parameters for k: [-103.02936496  102.82694212]
+  L2 norm of alpha & beta at the optimum: 2.1624999956742072e-05
+  Min. runtime: 7740.566666666667ns
+/home/axel/src/impactx/./match_beam.py:94: RuntimeWarning: Method CG cannot handle bounds.
+  res = minimize(
+CG 37
+  Optimal parameters for k: [-103.02936686  102.82694637]
+  L2 norm of alpha & beta at the optimum: 2.1624999902094705e-05
+  Min. runtime: 33754.16216216216ns
diff --git a/src/ImpactX.H b/src/ImpactX.H
index 274be58e7..57c5b3338 100644
--- a/src/ImpactX.H
+++ b/src/ImpactX.H
@@ -18,9 +18,27 @@
 #include <AMReX_REAL.H>
 
 #include <list>
+#include <map>
 #include <memory>
 
 
+// Enzyme
+//   https://enzyme.mit.edu/getting_started/Examples/
+//   https://enzyme.mit.edu/getting_started/CallingConvention/#types
+//   https://enzyme.mit.edu/getting_started/CallingConvention/#result-only-duplicated-argument
+inline int enzyme_dup;
+inline int enzyme_dupnoneed;
+inline int enzyme_in;
+inline int enzyme_out;
+inline int enzyme_const;
+inline int enzyme_width;
+
+amrex::ParticleReal __enzyme_fwddiff (...);
+amrex::ParticleReal __enzyme_autodiff (...);
+//
+// End Enzyme
+
+
 namespace impactx
 {
     /** An ImpactX simulation
@@ -161,6 +179,16 @@ namespace impactx
         bool m_grids_initialized = false;
     };
 
+    void
+    evaluate_lattice (ImpactX * sim, amrex::ParticleReal * error, amrex::ParticleReal const * q1_k, amrex::ParticleReal const * q2_k);
+
+    // optimal solution is:
+    //   q1_k = -103.12574100336;
+    //   q2_k = -q1_k;
+    // mode: "gradient-free", "forward", "backward"
+    std::map<std::string, amrex::ParticleReal>
+    my_run (amrex::ParticleReal q1_k = -85.0, amrex::ParticleReal q2_k=80.0, std::string mode="backward", std::string inputs_file="../../examples/fodo_space_charge/input_fodo_envelope_sc.in", bool verbose=true);
+
 } // namespace impactx
 
 #endif // IMPACT_X_H
diff --git a/src/ImpactX.cpp b/src/ImpactX.cpp
index e757b5318..447197f1e 100644
--- a/src/ImpactX.cpp
+++ b/src/ImpactX.cpp
@@ -11,17 +11,22 @@
 #include "initialization/InitAmrCore.H"
 #include "particles/ImpactXParticleContainer.H"
 #include "particles/Push.H"
+#include "elements/All.H"
+#include "envelope/spacecharge/EnvelopeSpaceChargePush.H"
 #include "diagnostics/DiagnosticOutput.H"
+#include "diagnostics/ReducedBeamCharacteristics.H"
 #include "particles/wakefields/HandleWakefield.H"
 
 #include <AMReX.H>
 #include <AMReX_AmrParGDB.H>
 #include <AMReX_BLProfiler.H>
+#include <AMReX_Math.H>
 #include <AMReX_ParallelDescriptor.H>
 #include <AMReX_ParmParse.H>
 #include <AMReX_Print.H>
 #include <AMReX_Utility.H>
 
+#include <array>
 #include <iostream>
 #include <memory>
 
@@ -149,4 +154,236 @@ namespace impactx {
         }
     }
 
+    void
+    evaluate_lattice (ImpactX * sim, amrex::ParticleReal * error, amrex::ParticleReal const * q1_k, amrex::ParticleReal const * q2_k) {
+        // Example from:
+        // https://impactx.readthedocs.io/en/latest/usage/examples/fodo_space_charge/README.html
+
+        // reference
+        // RefPart & ref = sim->amr_data->track_envelope.m_ref.value();
+        // auto & env = sim->amr_data->track_envelope.m_env.value();
+        // copy
+        RefPart ref = sim->amr_data->track_envelope.m_ref.value();
+        Envelope env = sim->amr_data->track_envelope.m_env.value();
+
+        int nslice = 50;
+
+        auto dr1 = elements::Drift{7.44e-2, 0, 0, 0, 0, 0, nslice};
+        auto q1 = elements::Quad{6.10e-2, *q1_k, 0, 0, 0, 0, 0, nslice};
+        auto dr2 = elements::Drift{14.88e-2, 0, 0, 0, 0, 0, nslice};
+        auto q2 = elements::Quad{6.10e-2, *q2_k, 0, 0, 0, 0, 0, nslice};
+
+        /*
+        // quadrupole triplet
+        // https://impactx.readthedocs.io/en/latest/usage/examples/optimize_triplet/README.html
+        sim->m_lattice = {
+            elements::Drift{2.7},
+            elements::Quad{0.1, q1_k},
+            elements::Drift{1.4},
+            elements::Quad{0.2, q2_k},
+            elements::Drift{1.4},
+            elements::Quad{0.1, q1_k},
+            elements::Drift{2.7}
+        };
+
+
+        // int ns = 25;  // number of slices per ds in the element
+        // for (auto & e : sim.m_lattice) { e.m_nslice = ns; }
+        */
+
+        //sim->evolve();
+        //auto & pc = *sim->amr_data->m_particle_container;
+        //dr1(pc, 0, 0);
+
+        // envelope mode
+        //std::cout << "before ref\n";
+        //std::cout << "before cm\n";
+        auto cm = env.m_env;
+
+        auto & intensity = env.m_beam_intensity;
+        intensity = 0.5;  // Ampere
+
+        // initial values alpha and beta
+        // goal: periodic solution where the resulting alpha and beta are the same as the initial ones.
+        //       the resulting -q1_k should be equal to q2_k (but we intentionally do not enforce that
+        //       during the simulation loop here, and leave that to the optimizer to find out).
+        auto rbc_goal = diagnostics::reduced_beam_characteristics(cm, ref);
+        amrex::ParticleReal const alpha_x_goal = rbc_goal["alpha_x"];
+        amrex::ParticleReal const alpha_y_goal = rbc_goal["alpha_y"];
+        amrex::ParticleReal const beta_x_goal = rbc_goal["beta_x"];
+        amrex::ParticleReal const beta_y_goal = rbc_goal["beta_y"];
+        std::array<amrex::ParticleReal, 4> alpha_beta_goal = {alpha_x_goal, alpha_y_goal, beta_x_goal, beta_y_goal};
+
+        // sub-steps for space charge within the element
+        for (int slice_step = 0; slice_step < nslice; ++slice_step)
+        {
+            envelope::spacecharge::space_charge2D_push(ref, cm, intensity, dr1.ds() / nslice);
+            dr1(ref);
+            dr1(cm, ref);
+        }
+
+        for (int slice_step = 0; slice_step < nslice; ++slice_step)
+        {
+            envelope::spacecharge::space_charge2D_push(ref, cm, intensity, q1.ds() / nslice);
+            q1(ref);
+            q1(cm, ref);
+        }
+
+        for (int slice_step = 0; slice_step < nslice; ++slice_step)
+        {
+            envelope::spacecharge::space_charge2D_push(ref, cm, intensity, dr2.ds() / nslice);
+            dr2(ref);
+            dr2(cm, ref);
+        }
+
+        for (int slice_step = 0; slice_step < nslice; ++slice_step)
+        {
+            envelope::spacecharge::space_charge2D_push(ref, cm, intensity, q2.ds() / nslice);
+            q2(ref);
+            q2(cm, ref);
+        }
+
+        for (int slice_step = 0; slice_step < nslice; ++slice_step)
+        {
+            envelope::spacecharge::space_charge2D_push(ref, cm, intensity, dr1.ds() / nslice);
+            dr1(ref);
+            dr1(cm, ref);
+        }
+
+        /*
+        // particles
+        std::unordered_map<std::string, amrex::ParticleReal> const rbc =
+                diagnostics::reduced_beam_characteristics(*active_sim->amr_data->m_particle_container);
+
+        return rbc.at("alpha_x");  // TOOD: alpha_x, alpha_y, beta_x, beta_y
+        */
+
+        // envelope
+        /*
+        std::unordered_map<std::string, amrex::ParticleReal> const rbc =
+            diagnostics::reduced_beam_characteristics(cm, ref);
+        return rbc.at("alpha_x");
+        */
+
+        auto rbc_is = diagnostics::reduced_beam_characteristics(cm, ref);
+        amrex::ParticleReal const alpha_x = rbc_is["alpha_x"];
+        amrex::ParticleReal const alpha_y = rbc_is["alpha_y"];
+        amrex::ParticleReal const beta_x = rbc_is["beta_x"];
+        amrex::ParticleReal const beta_y = rbc_is["beta_y"];
+        std::array<amrex::ParticleReal, 4> alpha_beta_is = {alpha_x, alpha_y, beta_x, beta_y};
+
+        *error =
+            amrex::Math::powi<2>(alpha_beta_is[0] - alpha_beta_goal[0]) +
+            amrex::Math::powi<2>(alpha_beta_is[1] - alpha_beta_goal[1]) +
+            amrex::Math::powi<2>(alpha_beta_is[2] - alpha_beta_goal[2]) +
+            amrex::Math::powi<2>(alpha_beta_is[3] - alpha_beta_goal[3]);
+    }
+
+    std::map<std::string, amrex::ParticleReal>
+    my_run (amrex::ParticleReal q1_k, amrex::ParticleReal q2_k, std::string mode, std::string inputs_file, bool verbose) {
+        amrex::ParmParse pp_ix("impactx");
+        pp_ix.add("verbose", 0);
+        amrex::ParmParse pp_amr("amrex");
+        pp_amr.add("verbose", 0);
+
+        ImpactX sim;
+
+        amrex::ParmParse pp_algo("algo");
+        std::string track = "envelope";
+        pp_algo.add("track", track);
+        pp_algo.add("particle_shape", 2);  // unused
+
+        sim.init_grids();
+
+        if (inputs_file != "") {
+            if (!amrex::FileExists(inputs_file)) {
+                throw std::runtime_error("my_run: inputs_file does not exist: " + inputs_file);
+            }
+            amrex::ParmParse::addfile(inputs_file);
+        }
+
+        // TODO: replace with beam params from https://impactx.readthedocs.io/en/latest/usage/examples/fodo_space_charge/README.html
+        sim.initBeamDistributionFromInputs();
+
+        // design the accelerator lattice
+        // sim.initLatticeElementsFromInputs();
+
+        // initial quad strengths
+        // optimal solution is:
+        //q1_k = -103.12574100336;
+        //q2_k = -q1_k;
+
+        amrex::ParticleReal error = 0.0;
+        amrex::ParticleReal derror_dq1_k = std::numeric_limits<amrex::ParticleReal>::quiet_NaN();
+        amrex::ParticleReal derror_dq2_k = std::numeric_limits<amrex::ParticleReal>::quiet_NaN();
+
+        //mode = "forward";
+        //mode = "backward";
+
+        if (mode == "gradient-free")
+        {
+            // non-differentiable run:
+            if (verbose) amrex::Print() << "Gradient-free run\n";
+            evaluate_lattice(&sim, &error, &q1_k, &q2_k);
+        }
+        else if (mode == "forward")
+        {
+            // forward differentiable run
+            if (verbose) amrex::Print() << "Forward differentiable run\n";
+
+            // in forward-mode, usually use one variable at a time, but enzyme_width allows batching
+            // two first order derivatives.
+            //   note: seeded direction, this is AD and NOT a finite difference
+            amrex::ParticleReal dq1_k = 1.0;  // derror += 1.0 * df/dq1_k
+            amrex::ParticleReal dq2_k = 1.0;  // derror += 1.0 * df/dq1_k
+            amrex::ParticleReal zero = 0.0;  // derror += 0.0 * df/dq*_k
+            __enzyme_fwddiff(
+                &evaluate_lattice,
+                enzyme_const, &sim,
+                enzyme_width, 2,  // 2x batched: first df/dq1_k then df/dq2_k
+                // variable, batch 1, batch 2:
+                enzyme_dup, &error, &derror_dq1_k, &derror_dq2_k,
+                enzyme_dup, &q1_k, &dq1_k, &zero,
+                enzyme_dup, &q2_k, &zero, &dq2_k
+            );
+        }
+        else if (mode == "backward")
+        {
+            // reverse differentiable run:
+            if (verbose) amrex::Print() << "Reverse differentiable run\n";
+            amrex::ParticleReal dq1_k = 0.0;  // accumulator, zero-init!
+            amrex::ParticleReal dq2_k = 0.0;  // accumulator, zero-init!
+            amrex::ParticleReal derror = 1.0;  // variable to back-prop: derror * df/dx
+            __enzyme_autodiff(
+                &evaluate_lattice,
+                enzyme_const, &sim,
+                enzyme_dup, &error, &derror,
+                enzyme_dup, &q1_k, &dq1_k,
+                enzyme_dup, &q2_k, &dq2_k
+            );
+            derror_dq1_k = dq1_k;
+            derror_dq2_k = dq2_k;
+        }
+        else {
+            throw std::runtime_error("Unknown mode: " + mode);
+        }
+
+        sim.finalize();
+
+        if (verbose) {
+            amrex::Print() << "q1_k = " << q1_k << "\n"
+                           << "q2_k = " << q2_k
+                           << std::endl;
+            amrex::Print() << "error = " << error << "\n"
+                           << "derror_dq1_k = " << derror_dq1_k << "\n"
+                           << "derror_dq2_k = " << derror_dq2_k << "\n"
+                           << std::endl;
+        }
+
+        std::map<std::string, amrex::ParticleReal> result;
+        result["error"] = error;
+        result["derror_dq1_k"] = derror_dq1_k;
+        result["derror_dq2_k"] = derror_dq2_k;
+        return result;
+    }
 } // namespace impactx
diff --git a/src/diagnostics/DiagnosticOutput.cpp b/src/diagnostics/DiagnosticOutput.cpp
index 3b5ca9682..22e317dab 100644
--- a/src/diagnostics/DiagnosticOutput.cpp
+++ b/src/diagnostics/DiagnosticOutput.cpp
@@ -198,8 +198,8 @@ namespace impactx::diagnostics
         prepare_header(file_handler, OutputType::PrintReducedBeamCharacteristics, append);
 
         amrex::ParticleReal const s = ref_part.s;
-        std::unordered_map<std::string, amrex::ParticleReal> const rbc =
-            diagnostics::reduced_beam_characteristics(cm, ref_part);
+        std::unordered_map<std::string, amrex::ParticleReal> const rbc; //=
+            //diagnostics::reduced_beam_characteristics(cm, ref_part);
 
         write_rbc(file_handler, rbc, s, step);
     }
diff --git a/src/diagnostics/ReducedBeamCharacteristics.H b/src/diagnostics/ReducedBeamCharacteristics.H
index da8aee6df..8d857b6ca 100644
--- a/src/diagnostics/ReducedBeamCharacteristics.H
+++ b/src/diagnostics/ReducedBeamCharacteristics.H
@@ -31,6 +31,7 @@ namespace impactx::diagnostics
     /** Compute moments of the beam distribution from covariance matrix
      */
     std::unordered_map<std::string, amrex::ParticleReal>
+    //amrex::ParticleReal
     reduced_beam_characteristics (Map6x6 const & cm, RefPart const & ref_part);
 
 } // namespace impactx::diagnostics
diff --git a/src/diagnostics/ReducedBeamCharacteristics.cpp b/src/diagnostics/ReducedBeamCharacteristics.cpp
index 7bcaa84e2..f5ea5a9c7 100644
--- a/src/diagnostics/ReducedBeamCharacteristics.cpp
+++ b/src/diagnostics/ReducedBeamCharacteristics.cpp
@@ -423,15 +423,16 @@ namespace impactx::diagnostics
     }
 
     std::unordered_map<std::string, amrex::ParticleReal>
+    //amrex::ParticleReal
     reduced_beam_characteristics (Map6x6 const & cm, RefPart const & ref_part)
     {
-        BL_PROFILE("impactx::diagnostics::reduced_beam_characteristics(cm)");
+        //BL_PROFILE("impactx::diagnostics::reduced_beam_characteristics(cm)");
 
         // reference particle relativistic beta*gamma
-        amrex::ParticleReal const bg = ref_part.beta_gamma();
-        amrex::ParticleReal const bg2 = bg*bg;
+        //amrex::ParticleReal const bg = ref_part.beta_gamma();
+        //amrex::ParticleReal const bg2 = bg*bg;
 
-       // mean square and correlation values
+        // mean square and correlation values
         amrex::ParticleReal const x_ms   = cm(1,1);
         amrex::ParticleReal const y_ms   = cm(3,3);
         amrex::ParticleReal const t_ms   = cm(5,5);
@@ -445,6 +446,7 @@ namespace impactx::diagnostics
         amrex::ParticleReal const pxpt   = cm(2,6);
         amrex::ParticleReal const ypt    = cm(3,6);
         amrex::ParticleReal const pypt   = cm(4,6);
+        /*
         amrex::ParticleReal const xy     = cm(1,3);
         amrex::ParticleReal const xpy    = cm(1,4);
         amrex::ParticleReal const xt     = cm(1,5);
@@ -453,20 +455,21 @@ namespace impactx::diagnostics
         amrex::ParticleReal const pxt    = cm(2,5);
         amrex::ParticleReal const yt     = cm(3,5);
         amrex::ParticleReal const pyt    = cm(4,5);
+        */
         // standard deviations of positions
-        amrex::ParticleReal const sig_x = std::sqrt(x_ms);
-        amrex::ParticleReal const sig_y = std::sqrt(y_ms);
-        amrex::ParticleReal const sig_t = std::sqrt(t_ms);
+        //amrex::ParticleReal const sig_x = std::sqrt(x_ms);
+        //amrex::ParticleReal const sig_y = std::sqrt(y_ms);
+        //amrex::ParticleReal const sig_t = std::sqrt(t_ms);
         // standard deviations of momenta
-        amrex::ParticleReal const sig_px = std::sqrt(px_ms);
-        amrex::ParticleReal const sig_py = std::sqrt(py_ms);
-        amrex::ParticleReal const sig_pt = std::sqrt(pt_ms);
+        //amrex::ParticleReal const sig_px = std::sqrt(px_ms);
+        //amrex::ParticleReal const sig_py = std::sqrt(py_ms);
+        //amrex::ParticleReal const sig_pt = std::sqrt(pt_ms);
         // RMS emittances
-        amrex::ParticleReal const e2_x = x_ms*px_ms-xpx*xpx;
-        amrex::ParticleReal const e2_y = y_ms*py_ms-ypy*ypy;
+        //amrex::ParticleReal const e2_x = x_ms*px_ms-xpx*xpx;
+        //amrex::ParticleReal const e2_y = y_ms*py_ms-ypy*ypy;
         amrex::ParticleReal const e2_t = t_ms*pt_ms-tpt*tpt;
-        amrex::ParticleReal const emittance_x = (e2_x > 0.0)? std::sqrt(e2_x) : 0.0;
-        amrex::ParticleReal const emittance_y = (e2_y > 0.0)? std::sqrt(e2_y) : 0.0;
+        //amrex::ParticleReal const emittance_x = (e2_x > 0.0)? std::sqrt(e2_x) : 0.0;
+        //amrex::ParticleReal const emittance_y = (e2_y > 0.0)? std::sqrt(e2_y) : 0.0;
         amrex::ParticleReal const emittance_t = (e2_t > 0.0)? std::sqrt(e2_t) : 0.0;
         // Dispersion and dispersive beam moments
         amrex::ParticleReal const dispersion_x = ((pt_ms > 0.0) ? (- xpt / pt_ms) : 0.0);
@@ -491,10 +494,11 @@ namespace impactx::diagnostics
         amrex::ParticleReal const alpha_t = - tpt / emittance_t;
 
         // Calculate normalized emittances
-        amrex::ParticleReal emittance_xn = emittance_x * bg;
-        amrex::ParticleReal emittance_yn = emittance_y * bg;
-        amrex::ParticleReal emittance_tn = emittance_t * bg;
+        //amrex::ParticleReal emittance_xn = emittance_x * bg;
+        //amrex::ParticleReal emittance_yn = emittance_y * bg;
+        //amrex::ParticleReal emittance_tn = emittance_t * bg;
 
+        /*
         // Determine whether to calculate eigenemittances, and initialize
         amrex::ParmParse pp_diag("diag");
         bool compute_eigenemittances = false;
@@ -548,10 +552,12 @@ namespace impactx::diagnostics
            emittance_2 = std::get<1>(emittances);
            emittance_3 = std::get<2>(emittances);
         }
+        */
 
-        auto const nan = std::numeric_limits<amrex::ParticleReal>::quiet_NaN();
+        //auto const nan = 0.0; // std::numeric_limits<amrex::ParticleReal>::quiet_NaN();
 
         std::unordered_map<std::string, amrex::ParticleReal> data;
+        /*
         data["x_mean"] = nan;
         data["x_min"] = nan;
         data["x_max"] = nan;
@@ -579,12 +585,14 @@ namespace impactx::diagnostics
         data["emittance_x"] = emittance_x;
         data["emittance_y"] = emittance_y;
         data["emittance_t"] = emittance_t;
+        */
         data["alpha_x"] = alpha_x;
         data["alpha_y"] = alpha_y;
         data["alpha_t"] = alpha_t;
         data["beta_x"] = beta_x;
         data["beta_y"] = beta_y;
         data["beta_t"] = beta_t;
+        /*
         data["dispersion_x"] = dispersion_x;
         data["dispersion_px"] = dispersion_px;
         data["dispersion_y"] = dispersion_y;
@@ -592,14 +600,19 @@ namespace impactx::diagnostics
         data["emittance_xn"] = emittance_xn;
         data["emittance_yn"] = emittance_yn;
         data["emittance_tn"] = emittance_tn;
+        */
+        /*
         if (compute_eigenemittances) {
            data["emittance_1"] = emittance_1;
            data["emittance_2"] = emittance_2;
            data["emittance_3"] = emittance_3;
         }
-        data["charge_C"] = nan;  // TODO: with space charge
+        */
+        //data["charge_C"] = nan;  // TODO: with space charge
 
         return data;
+        //return beta_x;
+        //return alpha_x;
     }
 
 } // namespace impactx::diagnostics
diff --git a/src/envelope/spacecharge/Elliptic_Integral.cpp b/src/envelope/spacecharge/Elliptic_Integral.cpp
index 84aa84531..d9917351f 100644
--- a/src/envelope/spacecharge/Elliptic_Integral.cpp
+++ b/src/envelope/spacecharge/Elliptic_Integral.cpp
@@ -74,10 +74,12 @@ namespace impactx::envelope::spacecharge
           uplim < y ||
           uplim < z )
         {
+          /*
           ablastr::warn_manager::WMRecordWarning(
               "algo.space_charge",
               "Invalid input arguments in function Elliptic_RD.",
               ablastr::warn_manager::WarnPriority::high);
+          */
           value = 0.0_prt;
         }
 
@@ -125,10 +127,12 @@ namespace impactx::envelope::spacecharge
           zn = ( zn + lamda ) * 0.25_prt;
         }
 
+        /*
         ablastr::warn_manager::WMRecordWarning(
             "algo.space_charge",
             "Exceeded maximum number of iterations in function Elliptic_RD.",
             ablastr::warn_manager::WarnPriority::medium);
+        */
         value = 0.0_prt;
         return value;
 
diff --git a/src/envelope/spacecharge/EnvelopeSpaceChargePush.cpp b/src/envelope/spacecharge/EnvelopeSpaceChargePush.cpp
index 519bcd3cf..26d78a765 100644
--- a/src/envelope/spacecharge/EnvelopeSpaceChargePush.cpp
+++ b/src/envelope/spacecharge/EnvelopeSpaceChargePush.cpp
@@ -111,6 +111,7 @@ namespace impactx::envelope::spacecharge
         amrex::ParticleReal const corr_xy = (xy==0.0)? 0.0 : std::abs(cm(1,3)/std::sqrt(xy));
         amrex::ParticleReal const corr_yz = (yz==0.0)? 0.0 : std::abs(cm(3,5)/std::sqrt(yz));
         amrex::ParticleReal const corr_zx = (zx==0.0)? 0.0 : std::abs(cm(5,1)/std::sqrt(zx));
+        /*
         if (corr_xy > errtol || corr_yz > errtol || corr_zx > errtol) {
             ablastr::warn_manager::WMRecordWarning(
                 "algo.space_charge",
@@ -119,6 +120,7 @@ namespace impactx::envelope::spacecharge
                 ablastr::warn_manager::WarnPriority::high
             );
         }
+        */
 
         // evaluate the off-identity elements of the linear transfer map
         R(2,1) = coeff * Elliptic_RD(y,z,x,errtol);
diff --git a/src/initialization/InitElement.cpp b/src/initialization/InitElement.cpp
index bd7dab8cd..d646d4ef2 100644
--- a/src/initialization/InitElement.cpp
+++ b/src/initialization/InitElement.cpp
@@ -567,7 +567,7 @@ element_name) );
                 pp_element.get("openpmd_path", openpmd_path);
             }
 
-            m_lattice.emplace_back( Source(distribution, openpmd_path, element_name) );
+            m_lattice.emplace_back(Source(distribution, openpmd_path, element_name) );
         } else if (element_type == "line")
         {
             // Parse the lattice elements for the sub-lattice in the line
diff --git a/src/main.cpp b/src/main.cpp
index aa3c973f8..d8f36ef8a 100644
--- a/src/main.cpp
+++ b/src/main.cpp
@@ -10,36 +10,17 @@
 #include "ImpactX.H"
 #include "initialization/InitAMReX.H"
 
-#include <AMReX.H>
-#include <AMReX_BLProfiler.H>
-
-#if defined(AMREX_USE_MPI)
-#   include <mpi.h>
-#endif
-
 
 int main(int argc, char* argv[])
 {
-#if defined(AMREX_USE_MPI)
-    AMREX_ALWAYS_ASSERT(MPI_SUCCESS == MPI_Init(&argc, &argv));
-#endif
+    using namespace impactx;
 
     // although ImpactX' init_grids will call this if not done before, we call
     // it here so users can pass command line arguments
-    impactx::initialization::default_init_AMReX(argc, argv);
+    initialization::default_init_AMReX(argc, argv);
 
     {
-        BL_PROFILE_VAR("main()", pmain);
-        impactx::ImpactX impactX;
-        impactX.init_grids();
-        impactX.initBeamDistributionFromInputs();
-        impactX.initLatticeElementsFromInputs();
-        impactX.evolve();
-        BL_PROFILE_VAR_STOP(pmain);
-        impactX.finalize();
+        my_run();
     }
 
-#if defined(AMREX_USE_MPI)
-    AMREX_ALWAYS_ASSERT(MPI_SUCCESS == MPI_Finalize());
-#endif
 }
diff --git a/src/python/pyImpactX.cpp b/src/python/pyImpactX.cpp
index 172b1f3ef..1c2030b13 100644
--- a/src/python/pyImpactX.cpp
+++ b/src/python/pyImpactX.cpp
@@ -53,6 +53,18 @@ PYBIND11_MODULE(impactx_pybind, m) {
     init_wakeconvolution(m);
     init_ImpactX(m);
 
+    m.def(
+        "my_run",
+        &my_run,
+        py::return_value_policy::copy,
+        py::arg("q1_k"),
+        py::arg("q2_k"),
+        py::arg("mode") = "backward",
+        py::arg("inputs_file") = "examples/fodo_space_charge/input_fodo_envelope_sc.in",
+        py::arg("verbose") = true,
+        "Run ImpactX lattice with mode: 'gradient-free', 'forward' or 'backward'"
+    );
+
     // expose our amrex module
     m.attr("amr") = amr;