Skip to content

Commit

Permalink
Use amrex::SmallMatrix
Browse files Browse the repository at this point in the history 
Replace current small matrix logic and math with the new
`amrex::SmallMatrix` class.
  • Loading branch information
@ax3l
ax3l committed Oct 10, 2024
1 parent 623cee4 commit 2eb6104
Show file tree
Hide file tree
Showing 8 changed files with 169 additions and 209 deletions.
2 changes: 1 addition & 1 deletion src/particles/ReferenceParticle.H
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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> map; ///< linearized map (zero-indexed)

/** Get reference particle relativistic gamma
*
Expand Down
54 changes: 1 addition & 53 deletions src/particles/diagnostics/CovarianceMatrixMath.H
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,17 +13,16 @@
#include <ablastr/constant.H>
#include <ablastr/warn_manager/WarnManager.H>

#include <AMReX_Array.H>
#include <AMReX_REAL.H>
#include <AMReX_GpuComplex.H>
#include <AMReX_SmallMatrix.H>

#include <cmath>
#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
Expand Down Expand Up @@ -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
14 changes: 8 additions & 6 deletions src/particles/diagnostics/EmittanceInvariants.H
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,10 +12,11 @@

#include "particles/ImpactXParticleContainer.H"

#include <AMReX_Array.H>
#include <AMReX_REAL.H>
#include <AMReX_SmallMatrix.H>

#include <string>
#include <tuple>



Expand Down Expand Up @@ -44,7 +45,7 @@ namespace impactx::diagnostics
amrex::ParticleReal
>
KineticInvariants (
amrex::Array2D<amrex::ParticleReal, 1, 6, 1, 6> const & Sigma
amrex::SmallMatrix<amrex::ParticleReal, 6, 6> const & Sigma
);

/** Returns the three eigenemittances
Expand All @@ -65,11 +66,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> const & Sigma
);

} // namespace impactx::diagnostics
Expand Down
56 changes: 28 additions & 28 deletions src/particles/diagnostics/EmittanceInvariants.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,54 +30,53 @@ 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> 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> S1;
amrex::SmallMatrix<amrex::ParticleReal, 6, 6> S2;
amrex::SmallMatrix<amrex::ParticleReal, 6, 6> S4;
amrex::SmallMatrix<amrex::ParticleReal, 6, 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++) {
for (int i = 0; i < 6; i++) {
for (int j = 0; j < 6; j++) {
if (j % 2 != 0) {
S1(i,j) = -Sigma(i,j+1); // if j is odd
S1(i,j) = +Sigma(i,j-1); // if j is odd
}
else {
S1(i,j) = +Sigma(i,j-1); // if j is even
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);
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);
return invariants;
Expand All @@ -95,11 +94,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> const & Sigma
)
{
BL_PROFILE("impactx::diagnostics::Eigenemittances");
Expand All @@ -123,8 +123,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";
Expand All @@ -133,10 +133,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;
Expand Down
82 changes: 42 additions & 40 deletions src/particles/diagnostics/ReducedBeamCharacteristics.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down Expand Up @@ -320,45 +321,46 @@ namespace impactx::diagnostics

if (compute_eigenemittances) {
// Store the covariance matrix in dynamical variables:
amrex::Array2D<amrex::ParticleReal, 1, 6, 1, 6> Sigma;
Sigma(1,1) = x_ms;
Sigma(1,2) = xpx * bg;
Sigma(1,3) = xy;
Sigma(1,4) = xpy * bg;
Sigma(1,5) = xt;
Sigma(1,6) = xpt * bg;
Sigma(2,1) = xpx * bg;
Sigma(2,2) = px_ms * bg2;
Sigma(2,3) = pxy * bg;
Sigma(2,4) = pxpy * bg2;
Sigma(2,5) = pxt * bg;
Sigma(2,6) = pxpt * bg2;
Sigma(3,1) = xy;
Sigma(3,2) = pxy * bg;
Sigma(3,3) = y_ms;
Sigma(3,4) = ypy * bg;
Sigma(3,5) = yt;
Sigma(3,6) = ypt * bg;
Sigma(4,1) = xpy * bg;
Sigma(4,2) = pxpy * bg2;
Sigma(4,3) = ypy * bg;
Sigma(4,4) = py_ms * bg2;
Sigma(4,5) = pyt * bg;
Sigma(4,6) = pypt * bg2;
Sigma(5,1) = xt;
Sigma(5,2) = pxt * bg;
Sigma(5,3) = yt;
Sigma(5,4) = pyt * bg;
Sigma(5,5) = t_ms;
Sigma(5,6) = tpt * bg;
Sigma(6,1) = xpt * bg;
Sigma(6,2) = pxpt * bg2;
Sigma(6,3) = ypt * bg;
Sigma(6,4) = pypt * bg2;
Sigma(6,5) = tpt * bg;
Sigma(6,6) = pt_ms * bg2;
// note: Sigma is zero-indexed
amrex::SmallMatrix<amrex::ParticleReal, 6, 6> Sigma;
Sigma(0 ,0) = x_ms;
Sigma(0, 1) = xpx * bg;
Sigma(0, 2) = xy;
Sigma(0, 3) = xpy * bg;
Sigma(0, 4) = xt;
Sigma(0, 5) = xpt * bg;
Sigma(1, 0) = xpx * bg;
Sigma(1, 1) = px_ms * bg2;
Sigma(1, 2) = pxy * bg;
Sigma(1, 3) = pxpy * bg2;
Sigma(1, 4) = pxt * bg;
Sigma(1, 5) = pxpt * bg2;
Sigma(2, 0) = xy;
Sigma(2, 1) = pxy * bg;
Sigma(2, 2) = y_ms;
Sigma(2, 3) = ypy * bg;
Sigma(2, 4) = yt;
Sigma(2, 5) = ypt * bg;
Sigma(3, 0) = xpy * bg;
Sigma(3, 1) = pxpy * bg2;
Sigma(3, 2) = ypy * bg;
Sigma(3, 3) = py_ms * bg2;
Sigma(3, 4) = pyt * bg;
Sigma(3, 5) = pypt * bg2;
Sigma(4, 0) = xt;
Sigma(4, 1) = pxt * bg;
Sigma(4, 2) = yt;
Sigma(4, 3) = pyt * bg;
Sigma(4, 4) = t_ms;
Sigma(4, 5) = tpt * bg;
Sigma(5, 0) = xpt * bg;
Sigma(5, 1) = pxpt * bg2;
Sigma(5, 2) = ypt * bg;
Sigma(5, 3) = pypt * bg2;
Sigma(5, 4) = tpt * bg;
Sigma(5, 5) = 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);
Expand Down
Loading

0 comments on commit 2eb6104

Please sign in to comment.