Use amrex::SmallMatrix

Replace current small matrix logic and math with the new
`amrex::SmallMatrix` class.

src/particles/ReferenceParticle.H

@@ -12,11 +12,11 @@
 
 #include <ablastr/constant.H>
 
-#include <AMReX_Array.H>
 #include <AMReX_BLassert.H>
 #include <AMReX_Extension.H>
 #include <AMReX_GpuQualifiers.H>
 #include <AMReX_REAL.H>
+#include <AMReX_SmallMatrix.H>
 
 #include <cmath>
 
@@ -41,7 +41,7 @@ namespace impactx
         amrex::ParticleReal charge = 0.0; ///< reference charge, in C
 
         amrex::ParticleReal sedge = 0.0;  ///< value of s at entrance of the current beamline element
-        amrex::Array2D<amrex::ParticleReal, 1, 6, 1, 6> map; ///< linearized map
+        amrex::SmallMatrix<amrex::ParticleReal, 6, 6, amrex::Order::F, 1> map{}; ///< linearized map
 
         /** Get reference particle relativistic gamma
          *
@@ -66,7 +66,7 @@ namespace impactx
             using namespace amrex::literals;
 
             amrex::ParticleReal const ref_gamma = -pt;
-            amrex::ParticleReal const ref_beta = std::sqrt(1.0_prt - 1.0_prt/pow(ref_gamma,2));
+            amrex::ParticleReal const ref_beta = std::sqrt(1.0_prt - 1.0_prt / std::pow(ref_gamma, 2));
             return ref_beta;
         }
 
@@ -95,7 +95,7 @@ namespace impactx
         {
             using namespace amrex::literals;
 
-            constexpr amrex::ParticleReal inv_MeV_invc2 = 1.0_prt /  ablastr::constant::SI::MeV_invc2;
+            constexpr amrex::ParticleReal inv_MeV_invc2 = 1.0_prt / ablastr::constant::SI::MeV_invc2;
             return amrex::ParticleReal(mass * inv_MeV_invc2);
         }
 

src/particles/diagnostics/CovarianceMatrixMath.H

@@ -13,7 +13,6 @@
 #include <ablastr/constant.H>
 #include <ablastr/warn_manager/WarnManager.H>
 
-#include <AMReX_Array.H>
 #include <AMReX_REAL.H>
 #include <AMReX_GpuComplex.H>
 
@@ -163,57 +162,6 @@ namespace impactx::diagnostics
         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> const & 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> const & A,
-        amrex::Array2D<amrex::ParticleReal, 1, 6, 1, 6> const & 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

@@ -12,11 +12,11 @@
 
 #include "particles/ImpactXParticleContainer.H"
 
-#include <AMReX_Array.H>
 #include <AMReX_REAL.H>
+#include <AMReX_SmallMatrix.H>
 
 #include <string>
-
+#include <tuple>
 
 
 namespace impactx::diagnostics
@@ -44,7 +44,7 @@ namespace impactx::diagnostics
         amrex::ParticleReal
     >
     KineticInvariants (
-        amrex::Array2D<amrex::ParticleReal, 1, 6, 1, 6> const & Sigma
+        amrex::SmallMatrix<amrex::ParticleReal, 6, 6, amrex::Order::F, 1> const & Sigma
     );
 
     /** Returns the three eigenemittances
@@ -65,11 +65,12 @@ namespace impactx::diagnostics
      * @returns tuple containing eigenemittances e1, e2, and e3
      */
     std::tuple<
-            amrex::ParticleReal,
-            amrex::ParticleReal,
-            amrex::ParticleReal>
+        amrex::ParticleReal,
+        amrex::ParticleReal,
+        amrex::ParticleReal
+    >
     Eigenemittances (
-        amrex::Array2D<amrex::ParticleReal, 1, 6, 1, 6> const & Sigma
+        amrex::SmallMatrix<amrex::ParticleReal, 6, 6, amrex::Order::F, 1> const & Sigma
     );
 
 } // namespace impactx::diagnostics

src/particles/diagnostics/EmittanceInvariants.cpp

@@ -7,12 +7,13 @@
  * Authors: Chad Mitchell, Axel Huebl
  * License: BSD-3-Clause-LBNL
  */
+#include "EmittanceInvariants.H"
 #include "CovarianceMatrixMath.H"
 
+#include <AMReX_BLProfiler.H>
 #include <AMReX_Extension.H>
-#include <AMReX_REAL.H>
 #include <AMReX_GpuComplex.H>
-#include <AMReX_BLProfiler.H>
+#include <AMReX_REAL.H>
 
 #include <cmath>
 #include <stdexcept>
@@ -30,29 +31,29 @@ namespace impactx::diagnostics
      *  calculation of the three eigenemittances.
      *
      * input - Sigma symmetric 6x6 covariance matrix
-     * returns - tuple containing invarants I2, I4, and I6
+     * 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> const & Sigma
+        amrex::SmallMatrix<amrex::ParticleReal, 6, 6, amrex::Order::F, 1> const & Sigma
     )
     {
         using namespace amrex::literals;
 
-        std::tuple <amrex::ParticleReal,amrex::ParticleReal,amrex::ParticleReal> invariants;
+        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;
+        amrex::SmallMatrix<amrex::ParticleReal, 6, 6, amrex::Order::F, 1> S1;
+        amrex::SmallMatrix<amrex::ParticleReal, 6, 6, amrex::Order::F, 1> S2;
+        amrex::SmallMatrix<amrex::ParticleReal, 6, 6, amrex::Order::F, 1> S4;
+        amrex::SmallMatrix<amrex::ParticleReal, 6, 6, amrex::Order::F, 1> S6;
 
         // Construct the matrix S1 = Sigma*J.  This is a
         // permutation of the columns of Sigma with
@@ -69,14 +70,14 @@ namespace impactx::diagnostics
         }
 
         // 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);
+        S2 = S1 * S1;
+        S4 = S2 * S2;
+        S6 = 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;
+        I2 = -S2.trace() / 2.0_prt;
+        I4 = +S4.trace() / 2.0_prt;
+        I6 = -S6.trace() / 2.0_prt;
 
 
         invariants = std::make_tuple(I2,I4,I6);
@@ -99,7 +100,7 @@ namespace impactx::diagnostics
             amrex::ParticleReal,
             amrex::ParticleReal>
     Eigenemittances (
-        amrex::Array2D<amrex::ParticleReal, 1, 6, 1, 6> const & Sigma
+        amrex::SmallMatrix<amrex::ParticleReal, 6, 6, amrex::Order::F, 1> const & Sigma
     )
     {
         BL_PROFILE("impactx::diagnostics::Eigenemittances");
@@ -123,8 +124,8 @@ namespace impactx::diagnostics
         // doi:10.48550/arXiv.1305.1532.
         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;
+        amrex::ParticleReal c = (std::pow(I2, 2) - I4) / 2.0_prt;
+        amrex::ParticleReal d = -std::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";
@@ -133,10 +134,10 @@ namespace impactx::diagnostics
         // Caution: The order of e1,e2,e3 should be consistent with the
         // order ex,ey,et in the limit of uncoupled transport.
         // The order below is important and has been checked.
-        roots = CubicRootsTrig(a,b,c,d);
-        amrex::ParticleReal e1 = sqrt(std::abs(std::get<1>(roots)));
-        amrex::ParticleReal e2 = sqrt(std::abs(std::get<2>(roots)));
-        amrex::ParticleReal e3 = sqrt(std::abs(std::get<0>(roots)));
+        roots = CubicRootsTrig(a, b, c, d);
+        amrex::ParticleReal e1 = std::sqrt(std::abs(std::get<1>(roots)));
+        amrex::ParticleReal e2 = std::sqrt(std::abs(std::get<2>(roots)));
+        amrex::ParticleReal e3 = std::sqrt(std::abs(std::get<0>(roots)));
 
         emittances = std::make_tuple(e1,e2,e3);
         return emittances;

src/particles/diagnostics/ReducedBeamCharacteristics.cpp

@@ -16,10 +16,11 @@
 
 #include <AMReX_BLProfiler.H>           // for TinyProfiler
 #include <AMReX_GpuQualifiers.H>        // for AMREX_GPU_DEVICE
-#include <AMReX_REAL.H>                 // for ParticleReal
-#include <AMReX_Reduce.H>               // for ReduceOps
 #include <AMReX_ParallelDescriptor.H>   // for ParallelDescriptor
 #include <AMReX_ParticleReduce.H>       // for ParticleReduce
+#include <AMReX_REAL.H>                 // for ParticleReal
+#include <AMReX_Reduce.H>               // for ReduceOps
+#include <AMReX_SmallMatrix.H>          // for SmallMatrix
 #include <AMReX_TypeList.H>             // for TypeMultiplier
 
 
@@ -320,7 +321,7 @@ namespace impactx::diagnostics
 
         if (compute_eigenemittances) {
            // Store the covariance matrix in dynamical variables:
-           amrex::Array2D<amrex::ParticleReal, 1, 6, 1, 6> Sigma;
+           amrex::SmallMatrix<amrex::ParticleReal, 6, 6, amrex::Order::F, 1> Sigma;
            Sigma(1,1) = x_ms;
            Sigma(1,2) = xpx * bg;
            Sigma(1,3) = xy;
@@ -358,7 +359,7 @@ namespace impactx::diagnostics
            Sigma(6,5) = tpt * bg;
            Sigma(6,6) = pt_ms * bg2;
            // Calculate eigenemittances
-           std::tuple <amrex::ParticleReal,amrex::ParticleReal,amrex::ParticleReal> emittances = Eigenemittances(Sigma);
+           std::tuple<amrex::ParticleReal, amrex::ParticleReal, amrex::ParticleReal> emittances = Eigenemittances(Sigma);
            emittance_1 = std::get<0>(emittances);
            emittance_2 = std::get<1>(emittances);
            emittance_3 = std::get<2>(emittances);

src/particles/elements/RFCavity.H

@@ -20,9 +20,9 @@
 #include <ablastr/constant.H>
 
 #include <AMReX.H>
-#include <AMReX_Array.H>
 #include <AMReX_Extension.H>
 #include <AMReX_REAL.H>
+#include <AMReX_SmallMatrix.H>
 
 #include <array>
 #include <cmath>
@@ -219,7 +219,7 @@ namespace RFCavityData
             amrex::ParticleReal ptout = pt;
 
             // get the linear map
-            amrex::Array2D<amrex::ParticleReal, 1, 6, 1, 6> const R = refpart.map;
+            amrex::SmallMatrix<amrex::ParticleReal, 6, 6, amrex::Order::F, 1> const R = refpart.map;
 
             // symplectic linear map for the RF cavity is computed using the
             // Hamiltonian formalism as described in:
@@ -228,18 +228,20 @@ namespace RFCavityData
             // so that, e.g., R(3,4) = dyf/dpyi.
 
             // push particles using the linear map
-            xout = R(1,1)*x + R(1,2)*px + R(1,3)*y
-                     + R(1,4)*py + R(1,5)*t + R(1,6)*pt;
-            pxout = R(2,1)*x + R(2,2)*px + R(2,3)*y
-                  + R(2,4)*py + R(2,5)*t + R(2,6)*pt;
-            yout = R(3,1)*x + R(3,2)*px + R(3,3)*y
-                     + R(3,4)*py + R(3,5)*t + R(3,6)*pt;
-            pyout = R(4,1)*x + R(4,2)*px + R(4,3)*y
-                  + R(4,4)*py + R(4,5)*t + R(4,6)*pt;
-            tout = R(5,1)*x + R(5,2)*px + R(5,3)*y
-                     + R(5,4)*py + R(5,5)*t + R(5,6)*pt;
-            ptout = R(6,1)*x + R(6,2)*px + R(6,3)*y
-                  + R(6,4)*py + R(6,5)*t + R(6,6)*pt;
+            // clang-format off
+            xout =  R(1,1)*x +  R(1,2)*px + R(1,3)*y
+                  + R(1,4)*py + R(1,5)*t +  R(1,6)*pt;
+            pxout = R(2,1)*x +  R(2,2)*px + R(2,3)*y
+                  + R(2,4)*py + R(2,5)*t +  R(2,6)*pt;
+            yout =  R(3,1)*x +  R(3,2)*px + R(3,3)*y
+                  + R(3,4)*py + R(3,5)*t +  R(3,6)*pt;
+            pyout = R(4,1)*x +  R(4,2)*px + R(4,3)*y
+                  + R(4,4)*py + R(4,5)*t +  R(4,6)*pt;
+            tout =  R(5,1)*x +  R(5,2)*px + R(5,3)*y
+                  + R(5,4)*py + R(5,5)*t +  R(5,6)*pt;
+            ptout = R(6,1)*x +  R(6,2)*px + R(6,3)*y
+                  + R(6,4)*py + R(6,5)*t +  R(6,6)*pt;
+            // clang-format on
 
             // assign updated values
             x = xout;
@@ -274,14 +276,7 @@ namespace RFCavityData
             amrex::ParticleReal const sedge = refpart.sedge;
 
             // initialize linear map (deviation) values
-            for (int i=1; i<7; i++) {
-               for (int j=1; j<7; j++) {
-                  if (i == j)
-                      refpart.map(i, j) = 1.0_prt;
-                  else
-                      refpart.map(i, j) = 0.0_prt;
-               }
-            }
+            refpart.map = decltype(refpart.map)::Identity();
 
             // length of the current slice
             amrex::ParticleReal const slice_ds = m_ds / nslice();
@@ -422,11 +417,11 @@ namespace RFCavityData
             }
 
             // push the linear map equations
-            amrex::Array2D<amrex::ParticleReal, 1, 6, 1, 6> const R = refpart.map;
+            amrex::SmallMatrix<amrex::ParticleReal, 6, 6, amrex::Order::F, 1> const R = refpart.map;
             amrex::ParticleReal const betgam = refpart.beta_gamma();
 
-            refpart.map(5,5) = R(5,5) + tau*R(6,5)/pow(betgam,3);
-            refpart.map(5,6) = R(5,6) + tau*R(6,6)/pow(betgam,3);
+            refpart.map(5,5) = R(5,5) + tau*R(6,5)/std::pow(betgam,3);
+            refpart.map(5,6) = R(5,6) + tau*R(6,6)/std::pow(betgam,3);
         }
 
         /** This pushes the reference particle and the linear map matrix
@@ -462,7 +457,7 @@ namespace RFCavityData
             refpart.pt = pt;
 
             // push the linear map equations
-            amrex::Array2D<amrex::ParticleReal, 1, 6, 1, 6> const R = refpart.map;
+            amrex::SmallMatrix<amrex::ParticleReal, 6, 6, amrex::Order::F, 1> const R = refpart.map;
             amrex::ParticleReal const s = tau/refpart.beta_gamma();
             amrex::ParticleReal const L = E0*ezp * std::sin(k*t+phi)/(2.0_prt*k);
 
@@ -514,7 +509,7 @@ namespace RFCavityData
             refpart.pt = pt - E0*(ezintf-ezint) * std::cos(k*t+phi);
 
             // push the linear map equations
-            amrex::Array2D<amrex::ParticleReal, 1, 6, 1, 6> const R = refpart.map;
+            amrex::SmallMatrix<amrex::ParticleReal, 6, 6, amrex::Order::F, 1> const R = refpart.map;
             amrex::ParticleReal const M = E0*(ezintf-ezint)*k * std::sin(k*t+phi);
             amrex::ParticleReal const L = E0*(ezpf-ezp) * std::sin(k*t+phi)/(2.0_prt*k)+M/2.0_prt;