Skip to content

Commit

Permalink
Use amrex::SmallMatrix (#736)
Browse files Browse the repository at this point in the history 
* Use `amrex::SmallMatrix`

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

* Matrix-Vector Multiplication

More readable dense operation.

* Eigenemittance: Fix Zero-Init

Ensure zero init of all non-set values
  • Loading branch information
@ax3l
ax3l authored Oct 12, 2024
1 parent 9355a57 commit 5f21458
Show file tree
Hide file tree
Showing 8 changed files with 84 additions and 218 deletions.
8 changes: 4 additions & 4 deletions src/particles/ReferenceParticle.H
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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>

Expand All @@ -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
*
Expand All @@ -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;
}

Expand Down Expand Up @@ -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);
}

Expand Down
52 changes: 0 additions & 52 deletions src/particles/diagnostics/CovarianceMatrixMath.H
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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>

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
15 changes: 8 additions & 7 deletions src/particles/diagnostics/EmittanceInvariants.H
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand Down
47 changes: 20 additions & 27 deletions src/particles/diagnostics/EmittanceInvariants.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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>
Expand All @@ -30,29 +31,21 @@ 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;
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{};

// Construct the matrix S1 = Sigma*J. This is a
// permutation of the columns of Sigma with
Expand All @@ -69,17 +62,17 @@ 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);
auto const S2 = S1 * S1;
auto const S4 = S2 * S2;
auto const 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;
amrex::ParticleReal const I2 = -S2.trace() / 2.0_prt;
amrex::ParticleReal const I4 = +S4.trace() / 2.0_prt;
amrex::ParticleReal const I6 = -S6.trace() / 2.0_prt;


invariants = std::make_tuple(I2,I4,I6);
std::tuple<amrex::ParticleReal, amrex::ParticleReal, amrex::ParticleReal> invariants = std::make_tuple(I2, I4, I6);
return invariants;
}

Expand All @@ -99,7 +92,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");
Expand All @@ -123,8 +116,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 +126,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
9 changes: 5 additions & 4 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,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;
Expand Down Expand Up @@ -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);
Expand Down
Loading

0 comments on commit 5f21458

Please sign in to comment.