[WIP] Clad Test

CMakeLists.txt

@@ -163,6 +163,9 @@ set(_ALL_TARGETS lib)
 set_target_properties(lib PROPERTIES
     POSITION_INDEPENDENT_CODE ON
     WINDOWS_EXPORT_ALL_SYMBOLS ON
+    UNITY_BUILD ON
+    UNITY_BUILD_MODE BATCH
+    UNITY_BUILD_BATCH_SIZE 1000
 )
 
 # executable application

README.md

@@ -1,154 +1,36 @@
-# ImpactX
+# Differentiable ImpactX
 
-[![CI Status](https://github.com/ECP-WarpX/impactx/actions/workflows/ubuntu.yml/badge.svg)](https://github.com/ECP-WarpX/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/)
 
-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
-
-[![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).
-
-## Developer Environment
-
-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
-
-Before you start, you will need a copy of the ImpactX source code:
+## Prepare Developer Environment
 
+### Once
 ```bash
-git clone [email protected]:ECP-WarpX/impactx.git
-cd impactx
+conda env create -n clad -c conda-forge clad clangdev clang
 ```
 
-## Compile
+### Always
 
 ```bash
-# find dependencies & configure
-cmake -S . -B build
+conda activate clad
 
-# compile
-cmake --build build -j 4
+export CC="clang"
+export CXX="clang++"
+export CXXFLAGS="-I ${CONDA_PREFIX}/include -fplugin=${CONDA_PREFIX}/lib/clad.so"
 ```
 
-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
 
+With the active developer env above:
 ```bash
-ccmake build
+cmake --fresh -S . -B build -DImpactX_MPI=OFF -DImpactX_COMPUTE=NOACC -DImpactX_OPENPMD=OFF
+cmake --build build -j 6
 ```
 
-or by adding arguments with `-D<OPTION>=<VALUE>` to the first CMake call, e.g.:
-
-```bash
-cmake -S . -B build -DImpactX_COMPUTE=CUDA -DImpactX_MPI=OFF
-```
 
-### Python Compile
+## Test/Run
 
+With the active developer env above:
 ```bash
-# find dependencies & configure
-cmake -S . -B build -DImpactX_PYTHON=ON
-
-# compile & install
-cmake --build build -j 4 --target pip_install
-```
-
-## 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
-
-  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 voerwrite 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 wheel pytest
-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
+./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`.

examples/fodo/input_fodo.in

@@ -18,33 +18,6 @@ beam.muypy = -beam.muxpx
 beam.mutpt = 0.0
 
 
-###############################################################################
-# Beamline: lattice elements and segments
-###############################################################################
-lattice.elements = monitor drift1 monitor quad1 monitor drift2 monitor quad2 monitor drift3 monitor
-lattice.nslice = 25
-
-monitor.type = beam_monitor
-monitor.backend = h5
-
-drift1.type = drift
-drift1.ds = 0.25
-
-quad1.type = quad
-quad1.ds = 1.0
-quad1.k = 1.0
-
-drift2.type = drift
-drift2.ds = 0.5
-
-quad2.type = quad
-quad2.ds = 1.0
-quad2.k = -1.0
-
-drift3.type = drift
-drift3.ds = 0.25
-
-
 ###############################################################################
 # Algorithms
 ###############################################################################

src/ImpactX.H

@@ -17,6 +17,8 @@
 
 #include <AMReX_REAL.H>
 
+#include <clad/Differentiator/Differentiator.h>
+
 #include <list>
 #include <memory>
 
@@ -147,6 +149,14 @@ namespace impactx
         bool m_grids_initialized = false;
     };
 
+    amrex::ParticleReal
+    evaluate_lattice (amrex::ParticleReal q1_k, amrex::ParticleReal q2_k);
+
+    void my_run ();
+
+    // global var until functors work better
+    static ImpactX* active_sim = nullptr;
+
 } // namespace impactx
 
 #endif // IMPACT_X_H

src/ImpactX.cpp

@@ -13,7 +13,9 @@
 #include "particles/CollectLost.H"
 #include "particles/ImpactXParticleContainer.H"
 #include "particles/Push.H"
+#include "particles/elements/All.H"
 #include "particles/diagnostics/DiagnosticOutput.H"
+#include "particles/diagnostics/ReducedBeamCharacteristics.H"
 #include "particles/spacecharge/ForceFromSelfFields.H"
 #include "particles/spacecharge/GatherAndPush.H"
 #include "particles/spacecharge/PoissonSolve.H"
@@ -336,4 +338,78 @@ namespace impactx {
             }, element_variant);
         }
     }
+
+    amrex::ParticleReal
+    evaluate_lattice (amrex::ParticleReal q1_k, amrex::ParticleReal q2_k)
+    {
+        // ns = 10  // TODO: number of slices per ds in the element
+
+        auto dr1 = Drift{2.7};
+        auto q1 = Quad{0.1, q1_k};
+        auto dr2 = Drift{1.4};
+        auto q2 = Quad{0.2, q2_k};
+        auto dr3 = Drift{1.4};
+        auto q3 = Quad{0.1, q1_k};
+        auto dr4 = Drift{2.7};
+        /*
+        // quadrupole triplet
+        // https://impactx.readthedocs.io/en/latest/usage/examples/optimize_triplet/README.html
+        active_sim->m_lattice = {
+            Drift{2.7},
+            Quad{0.1, q1_k},
+            Drift{1.4},
+            Quad{0.2, q2_k},
+            Drift{1.4},
+            Quad{0.1, q1_k},
+            Drift{2.7}
+        };
+
+
+        // int ns = 25;  // number of slices per ds in the element
+        // for (auto & e : sim.m_lattice) { e.m_nslice = ns; }
+        */
+
+        //active_sim->evolve();
+        auto & pc = *active_sim->amr_data->m_particle_container;
+
+        dr1(pc, 0, 0);
+
+        /*
+        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
+        */
+        return 0.0;
+    }
+
+    void my_run ()
+    {
+        ImpactX sim;
+
+        sim.init_grids();
+
+        // TODO: replace with beam params from https://impactx.readthedocs.io/en/latest/usage/examples/optimize_triplet/README.html
+        sim.initBeamDistributionFromInputs();
+
+        // design the accelerator lattice
+        // sim.initLatticeElementsFromInputs();
+
+        // initial quad strengths
+        amrex::ParticleReal q1_k = -3.0;
+        amrex::ParticleReal q2_k = 3.0;
+
+        active_sim = &sim;
+
+        // non-differentiable run:
+        amrex::ParticleReal const alpha_x = evaluate_lattice(q1_k, q2_k);
+        amrex::Print() << "final alpha_x = " << alpha_x << std::endl;
+
+        // differentiable run:
+        auto fn_dx = clad::differentiate(evaluate_lattice, "q1_k");  // TODO: this fails with a Clad ICE
+        //fn_dx.dump();
+        //fn_dx.execute(5, 3);
+
+        sim.finalize();
+    }
 } // namespace impactx

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
 }