Add beam eigenemittances to reduced diagnostics.

src/particles/diagnostics/CMakeLists.txt

@@ -2,4 +2,5 @@ target_sources(lib
   PRIVATE
     ReducedBeamCharacteristics.cpp
     DiagnosticOutput.cpp
+    EmittanceInvariants.cpp
 )

src/particles/diagnostics/CovarianceMatrixMath.H

@@ -0,0 +1,137 @@
+/* Copyright 2022-2023 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.
+ *
+ * This file is part of ImpactX.
+ *
+ * Authors: Chad Mitchell, Axel Huebl
+ * License: BSD-3-Clause-LBNL
+ */
+#ifndef COVARIANCE_MATRIX_MATH_H
+#define COVARIANCE_MATRIX_MATH_H
+
+#include <ablastr/constant.H>
+#include <AMReX_Array.H>
+#include <AMReX_REAL.H>
+#include <tuple>
+
+namespace impactx::diagnostics
+{
+
+    /** Function to return the roots of a cubic polynomial ax^3 + bx^2 + cx + d.
+     *  The trigonometric form of Cardano's formula is used.
+     *  This implementation expects three real roots, which is verified
+     *  by checking the sign of the discriminant.
+     *
+     * @param[in] a coefficient of cubic term
+     * @param[in] b coefficient of quadratic term
+     * @param[in] c coefficient of linear term
+     * @param[in] d coefficient of constant term
+     * @returns tuple of three real roots
+     */
+    std::tuple<
+            amrex::ParticleReal,
+            amrex::ParticleReal,
+            amrex::ParticleReal>
+    CubicRoots (
+        amrex::ParticleReal a,
+        amrex::ParticleReal b,
+        amrex::ParticleReal c,
+        amrex::ParticleReal d
+    )
+    {
+        using namespace amrex::literals;
+        using ablastr::constant::math::pi;
+
+        std::tuple <amrex::ParticleReal,amrex::ParticleReal,amrex::ParticleReal> roots;
+        amrex::ParticleReal x1 = 0.0_prt;
+        amrex::ParticleReal x2 = 0.0_prt;
+        amrex::ParticleReal x3 = 0.0_prt;
+
+        amrex::ParticleReal Q = (3.0_prt*a*c - pow(b,2))/(9.0_prt * pow(a,2));
+        amrex::ParticleReal R = (9.0_prt*a*b*c - 27_prt*pow(a,2)*d - 2.0_prt*pow(b,3))/(54.0_prt*pow(a,3));
+        amrex::ParticleReal discriminant = pow(Q,3) + pow(R,2);
+
+        // Discriminant should be < 0.  Otherwise, keep theta at default and throw an error.
+        amrex::ParticleReal tol = 1.0e-12;  //allow for roundoff error
+        if(discriminant > tol){
+
+           std::cout << "Polynomial in CubicRoots has one or more complex roots." << "\n";
+
+        } else if (Q == 0.0_prt){  // Special case of a triple root
+
+           x1 = - b/(3.0_prt*a);
+           x2 = - b/(3.0_prt*a);
+           x3 = - b/(3.0_prt*a);
+
+        } else {
+
+           //Three real roots in trigonometric form.
+           amrex::ParticleReal theta = acos(R/sqrt(-pow(Q,3)));
+           x1 = 2.0_prt*sqrt(-Q)*cos(theta/3.0_prt) - b/(3.0_prt*a);
+           x2 = 2.0_prt*sqrt(-Q)*cos(theta/3.0_prt + 2.0_prt*pi/3.0_prt) - b/(3.0_prt*a);
+           x3 = 2.0_prt*sqrt(-Q)*cos(theta/3.0_prt + 4.0_prt*pi/3.0_prt) - b/(3.0_prt*a);
+
+        }
+
+        //std::cout << "Discriminant, Q, R " << discriminant << " " << Q << " " << R << "\n";
+        //std::cout << "Return x1, x2, x3 " << x1 << " " << x2 << " " << x3 << "\n";
+
+        roots = std::make_tuple(x1,x2,x3);
+        return roots;
+    }
+
+
+    /** Function to take the trace of a square 6x6 matrix.
+     *
+     * @param[in] A a square matrix
+     * @returns the trace of A
+     */
+    amrex::ParticleReal
+    TraceMat (
+        amrex::Array2D<amrex::ParticleReal, 1, 6, 1, 6> A
+    )
+    {
+        int const dim = 6;
+        amrex::ParticleReal trA = 0.0;
+
+        for (int i = 1; i < dim+1; i++) {
+            trA += A(i,i);
+        }
+        return trA;
+    }
+
+
+    /** Function to multiply two square matrices of dimension 6.
+     *
+     * @param[in] A a square matrix
+     * @param[in] B square matrix
+     * @returns the matrix C = AB
+     */
+    amrex::Array2D<amrex::ParticleReal, 1, 6, 1, 6>
+    MultiplyMat (
+        amrex::Array2D<amrex::ParticleReal, 1, 6, 1, 6> A,
+        amrex::Array2D<amrex::ParticleReal, 1, 6, 1, 6> B
+    )
+    {
+        amrex::Array2D<amrex::ParticleReal, 1, 6, 1, 6> C;
+        int const dim = 6;
+
+        for (int i = 1; i < dim+1; i++) {
+            for (int j = 1; j < dim+1; j++) {
+                C(i,j) = 0;
+
+                for (int k = 1; k < dim+1; k++) {
+                    C(i,j) += A(i,k) * B(k,j);
+                }
+
+            }
+
+        }
+        return C;
+    }
+
+
+} // namespace impactx::diagnostics
+
+#endif // COVARIANCE_MATRIX_MATH_H

src/particles/diagnostics/EmittanceInvariants.H

@@ -0,0 +1,61 @@
+/* Copyright 2022-2023 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.
+ *
+ * This file is part of ImpactX.
+ *
+ * Authors: Marco Garten, Chad Mitchell, Axel Huebl
+ * License: BSD-3-Clause-LBNL
+ */
+#ifndef IMPACTX_EMITTANCE_INVARIANTS
+#define IMPACTX_EMITTANCE_INVARIANTS
+
+#include "particles/ImpactXParticleContainer.H"
+
+#include <AMReX_REAL.H>
+
+#include <string>
+#include <unordered_map>
+
+
+namespace impactx::diagnostics
+{
+
+    /** Returns the three independent kinetic invariants
+     *  denoted I2, I4, and I6 as constructed from the 6x6
+     *  beam covariance matrix.  These three quantities are invariant
+     *  under any linear symplectic transport map, and are used in the
+     *  calculation of the three eigenemittances.
+     *
+     * @param[in] Sigma symmetric 6x6 covariance matrix
+     * @returns tuple containing invariants I2, I4, and I6
+     */
+    std::tuple<
+            amrex::ParticleReal,
+            amrex::ParticleReal,
+            amrex::ParticleReal>
+    KineticInvariants (
+        amrex::Array2D<amrex::ParticleReal, 1, 6, 1, 6> Sigma
+    );
+
+    /** Returns the three eigenemittances
+     *  denoted e1, e2, and e3 as constructed from the 6x6
+     *  beam covariance matrix.  These three quantities are invariant
+     *  under any linear symplectic transport map, and reduce to
+     *  the projected normalized rms emittances in the limit of
+     *  uncoupled transport.
+     *
+     * @param[in] Sigma symmetric 6x6 covariance matrix
+     * @returns tuple containing eigenemittances e1, e2, and e3
+     */
+    std::tuple<
+            amrex::ParticleReal,
+            amrex::ParticleReal,
+            amrex::ParticleReal>
+    Eigenemittances (
+        amrex::Array2D<amrex::ParticleReal, 1, 6, 1, 6> Sigma
+    );
+
+} // namespace impactx::diagnostics
+
+#endif // IMPACTX_EMITTANCE_INVARIANTS

src/particles/diagnostics/EmittanceInvariants.cpp

@@ -0,0 +1,138 @@
+/* Copyright 2022-2023 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.
+ *
+ * This file is part of ImpactX.
+ *
+ * Authors: Chad Mitchell, Axel Huebl
+ * License: BSD-3-Clause-LBNL
+ */
+#include "CovarianceMatrixMath.H"
+
+#include <stdexcept>
+#include <AMReX_Extension.H>
+#include <AMReX_REAL.H>
+#include <AMReX_GpuComplex.H>
+#include <cmath>
+#include <vector>
+#include <tuple>
+
+
+namespace impactx::diagnostics
+{
+
+    /** This function returns the three independent kinetic invariants
+     *  denoted I2, I4, and I6 as constructed from the 6x6
+     *  beam covariance matrix.  These three quantities are invariant
+     *  under any linear symplectic transport map, and are used in the
+     *  calculation of the three eigenemittances.
+     *
+     * input - Sigma symmetric 6x6 covariance matrix
+     * returns - tuple containing invarants I2, I4, and I6
+     */
+    std::tuple<
+            amrex::ParticleReal,
+            amrex::ParticleReal,
+            amrex::ParticleReal>
+    KineticInvariants (
+        amrex::Array2D<amrex::ParticleReal, 1, 6, 1, 6> Sigma
+    )
+    {
+        using namespace amrex::literals;
+
+        std::tuple <amrex::ParticleReal,amrex::ParticleReal,amrex::ParticleReal> invariants;
+        amrex::ParticleReal I2 = 0.0_prt;
+        amrex::ParticleReal I4 = 0.0_prt;
+        amrex::ParticleReal I6 = 0.0_prt;
+
+        // Intermediate matrices used for storage.
+        amrex::Array2D<amrex::ParticleReal, 1, 6, 1, 6> S1;
+        amrex::Array2D<amrex::ParticleReal, 1, 6, 1, 6> S2;
+        amrex::Array2D<amrex::ParticleReal, 1, 6, 1, 6> S4;
+        amrex::Array2D<amrex::ParticleReal, 1, 6, 1, 6> S6;
+
+        // Construct the matrix S1 = Sigma*J.  This is a
+        // permutation of the columns of Sigma with
+        // a change of sign.
+        for (int i = 1; i < 7; i++) {
+            for (int j = 1; j < 7; j++) {
+                if (j % 2) {
+                   S1(i,j) = -Sigma(i,j+1); // if j is odd
+                }
+                else {
+                   S1(i,j) = +Sigma(i,j-1); // if j is even
+                }
+            }
+        }
+
+        // Carry out necessary matrix multiplications (3 are needed).
+        S2 = impactx::diagnostics::MultiplyMat(S1,S1);
+        S4 = impactx::diagnostics::MultiplyMat(S2,S2);
+        S6 = impactx::diagnostics::MultiplyMat(S2,S4);
+
+        // Define the three kinematic invariants (should be nonnegative).
+        I2 = -impactx::diagnostics::TraceMat(S2)/2.0_prt;
+        I4 = +impactx::diagnostics::TraceMat(S4)/2.0_prt;
+        I6 = -impactx::diagnostics::TraceMat(S6)/2.0_prt;
+
+        //std::cout << "Return I2, I4, I6 " << I2 << " " << I4 << " " << I6 << "\n";
+
+        invariants = std::make_tuple(I2,I4,I6);
+        return invariants;
+    }
+
+
+    /** This function returns the three eigenemittances
+     *  denoted e1, e2, and e3 as constructed from the 6x6
+     *  beam covariance matrix.  These three quantities are invariant
+     *  under any linear symplectic transport map, and reduce to
+     *  the projected normalized rms emittances in the limit of
+     *  uncoupled transport.
+     *
+     * input - Sigma symmetric 6x6 covariance matrix
+     * returns - tuple containing eigenemittances e1, e2, and e3
+     */
+    std::tuple<
+            amrex::ParticleReal,
+            amrex::ParticleReal,
+            amrex::ParticleReal>
+    Eigenemittances (
+        amrex::Array2D<amrex::ParticleReal, 1, 6, 1, 6> Sigma
+    )
+    {
+        using namespace amrex::literals;
+
+        std::tuple <amrex::ParticleReal,amrex::ParticleReal,amrex::ParticleReal> invariants;
+        std::tuple <amrex::ParticleReal,amrex::ParticleReal,amrex::ParticleReal> roots;
+        std::tuple <amrex::ParticleReal,amrex::ParticleReal,amrex::ParticleReal> emittances;
+
+        // Get the invariants I2, I4, and I6 from the covariance matrix.
+        invariants = KineticInvariants(Sigma);
+        amrex::ParticleReal I2 = std::get<0>(invariants);
+        amrex::ParticleReal I4 = std::get<1>(invariants);
+        amrex::ParticleReal I6 = std::get<2>(invariants);
+
+        // Construct the coefficients of the cubic polynomial
+        amrex::ParticleReal a = 1.0_prt;
+        amrex::ParticleReal b = -I2;
+        amrex::ParticleReal c = (pow(I2,2)-I4)/2.0_prt;
+        amrex::ParticleReal d = -pow(I2,3)/6.0_prt + I2*I4/2.0_prt - I6/3.0_prt;
+
+        // Return the cubic coefficients
+        //std::cout << "Return a,b,c,d " << a << " " << b << " " << c << " " << d << "\n";
+
+        // Solve for the roots to obtain the eigenemittances.
+        roots = CubicRoots(a,b,c,d);
+        amrex::ParticleReal e1 = sqrt(std::get<1>(roots));
+        amrex::ParticleReal e2 = sqrt(std::get<2>(roots));
+        amrex::ParticleReal e3 = sqrt(std::get<0>(roots));
+
+        // Caution: The order of e1,e2,e3 should be consistent with the
+        // order ex,ey,et in the limit of uncoupled transport.  The
+        // ordering remains to be carefully checked (TODO).
+        emittances = std::make_tuple(e1,e2,e3);
+        return emittances;
+    }
+
+
+} // namespace impactx::diagnostics