Skip to content

Commit 9c39e1f

Browse files
cemitch99ax3l
cemitch99
and
ax3l
committed
Add covariance map diagnostics (#4)
* Update DiagnosticOutput & Reduced Beam Diagnostics --------- Co-authored-by: Axel Huebl <axel.huebl@plasma.ninja>
1 parent 53c5348 commit 9c39e1f

File tree

4 files changed

+379
-137
lines changed

4 files changed

+379
-137
lines changed

‎src/particles/diagnostics/DiagnosticOutput.H

Lines changed: 30 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
#define IMPACTX_DIAGNOSTIC_OUTPUT_H
1212

1313
#include "particles/ImpactXParticleContainer.H"
14+
#include "particles/CovarianceMatrix.H"
1415

1516
#include <string>
1617

@@ -38,11 +39,35 @@ namespace impactx::diagnostics
3839
* @param step the global step
3940
* @param append open a new file with a fresh header (false) or append data to an existing file (true)
4041
*/
41-
void DiagnosticOutput (ImpactXParticleContainer const & pc,
42-
OutputType otype,
43-
std::string file_name,
44-
int step = 0,
45-
bool append = false);
42+
void DiagnosticOutput (
43+
ImpactXParticleContainer const & pc,
44+
OutputType const otype,
45+
std::string file_name,
46+
int step = 0,
47+
bool append = false
48+
);
49+
50+
/** ASCII output diagnostics associated with the beam.
51+
*
52+
* This temporary implementation uses ASCII output.
53+
* It is intended only for small tests where IO performance is not
54+
* a concern. The implementation here serializes IO.
55+
*
56+
* @param cm covariance matrix
57+
* @param ref_part reference particle
58+
* @param otype the type of output to produce
59+
* @param file_name the file name to write to
60+
* @param step the global step
61+
* @param append open a new file with a fresh header (false) or append data to an existing file (true)
62+
*/
63+
void DiagnosticOutput (
64+
Map6x6 const & cm,
65+
RefPart const & ref_part,
66+
OutputType const otype,
67+
std::string file_name,
68+
int step = 0,
69+
bool append = false
70+
);
4671

4772
} // namespace impactx::diagnostics
4873

‎src/particles/diagnostics/DiagnosticOutput.cpp

Lines changed: 182 additions & 130 deletions
Original file line numberDiff line numberDiff line change
@@ -9,162 +9,179 @@
99
*/
1010
#include "DiagnosticOutput.H"
1111
#include "NonlinearLensInvariants.H"
12+
#include "particles/CovarianceMatrix.H"
1213
#include "ReducedBeamCharacteristics.H"
1314

1415
#include <AMReX_BLProfiler.H> // for BL_PROFILE
15-
#include <AMReX_Extension.H> // for AMREX_RESTRICT
1616
#include <AMReX_ParmParse.H> // for ParmParse
1717
#include <AMReX_REAL.H> // for ParticleReal
1818
#include <AMReX_Print.H> // for PrintToFile
19-
#include <AMReX_ParticleTile.H> // for constructor of SoAParticle
2019

2120
#include <limits>
21+
#include <stdexcept>
2222
#include <utility>
2323

2424

25-
namespace impactx::diagnostics
25+
namespace
2626
{
27-
void DiagnosticOutput (ImpactXParticleContainer const & pc,
28-
OutputType const otype,
29-
std::string file_name,
30-
int step,
31-
bool append)
27+
void
28+
write_column_header (
29+
amrex::AllPrintToFile & file_handler,
30+
impactx::diagnostics::OutputType const otype,
31+
bool has_charge
32+
)
3233
{
33-
BL_PROFILE("impactx::diagnostics::DiagnosticOutput");
34+
if (otype == impactx::diagnostics::OutputType::PrintRefParticle)
35+
{
36+
file_handler << "step s beta gamma beta_gamma x y z t px py pz pt\n";
37+
}
38+
else if (otype == impactx::diagnostics::OutputType::PrintReducedBeamCharacteristics)
39+
{
40+
// determine whether to output eigenemittances
41+
amrex::ParmParse pp_diag("diag");
42+
bool compute_eigenemittances = false;
43+
pp_diag.queryAdd("eigenemittances", compute_eigenemittances);
3444

35-
using namespace amrex::literals; // for _rt and _prt
45+
file_handler << "step" << " " << "s" << " "
46+
<< "x_mean" << " " << "x_min" << " " << "x_max" << " "
47+
<< "y_mean" << " " << "y_min" << " " << "y_max" << " "
48+
<< "t_mean" << " " << "t_min" << " " << "t_max" << " "
49+
<< "sig_x" << " " << "sig_y" << " " << "sig_t" << " "
50+
<< "px_mean" << " " << "px_min" << " " << "px_max" << " "
51+
<< "py_mean" << " " << "py_min" << " " << "py_max" << " "
52+
<< "pt_mean" << " " << "pt_min" << " " << "pt_max" << " "
53+
<< "sig_px" << " " << "sig_py" << " " << "sig_pt" << " "
54+
<< "emittance_x" << " " << "emittance_y" << " " << "emittance_t" << " "
55+
<< "alpha_x" << " " << "alpha_y" << " " << "alpha_t" << " "
56+
<< "beta_x" << " " << "beta_y" << " " << "beta_t" << " "
57+
<< "dispersion_x" << " " << "dispersion_px" << " "
58+
<< "dispersion_y" << " " << "dispersion_py" << " "
59+
<< "emittance_xn" << " " << "emittance_yn" << " " << "emittance_tn";
60+
if (compute_eigenemittances)
61+
{
62+
file_handler << " "
63+
<< "emittance_xn" << " " << "emittance_yn" << " " << "emittance_tn";
64+
}
65+
if (has_charge)
66+
{
67+
file_handler << " " << "charge_C";
68+
}
69+
file_handler << "\n";
70+
}
71+
}
3672

37-
// keep file open as we add more and more lines
38-
amrex::AllPrintToFile file_handler(std::move(file_name));
73+
void
74+
prepare_header (
75+
amrex::AllPrintToFile & file_handler,
76+
impactx::diagnostics::OutputType const otype,
77+
bool has_charge,
78+
bool append
79+
)
80+
{
3981
file_handler.SetPrecision(std::numeric_limits<amrex::ParticleReal>::max_digits10);
4082

4183
// write file header per MPI RANK
42-
if (!append) {
43-
if (otype == OutputType::PrintRefParticle) {
44-
file_handler << "step s beta gamma beta_gamma x y z t px py pz pt\n";
45-
} else if (otype == OutputType::PrintReducedBeamCharacteristics) {
46-
47-
// determine whether to output eigenemittances
48-
amrex::ParmParse pp_diag("diag");
49-
bool compute_eigenemittances = false;
50-
pp_diag.queryAdd("eigenemittances", compute_eigenemittances);
51-
52-
if (compute_eigenemittances) {
53-
file_handler << "step" << " " << "s" << " "
54-
<< "x_mean" << " " << "x_min" << " " << "x_max" << " "
55-
<< "y_mean" << " " << "y_min" << " " << "y_max" << " "
56-
<< "t_mean" << " " << "t_min" << " " << "t_max" << " "
57-
<< "sig_x" << " " << "sig_y" << " " << "sig_t" << " "
58-
<< "px_mean" << " " << "px_min" << " " << "px_max" << " "
59-
<< "py_mean" << " " << "py_min" << " " << "py_max" << " "
60-
<< "pt_mean" << " " << "pt_min" << " " << "pt_max" << " "
61-
<< "sig_px" << " " << "sig_py" << " " << "sig_pt" << " "
62-
<< "emittance_x" << " " << "emittance_y" << " " << "emittance_t" << " "
63-
<< "alpha_x" << " " << "alpha_y" << " " << "alpha_t" << " "
64-
<< "beta_x" << " " << "beta_y" << " " << "beta_t" << " "
65-
<< "dispersion_x" << " " << "dispersion_px" << " "
66-
<< "dispersion_y" << " " << "dispersion_py" << " "
67-
<< "emittance_xn" << " " << "emittance_yn" << " " << "emittance_tn" << " "
68-
<< "emittance_1" << " " << "emittance_2" << " " << "emittance_3" << " "
69-
<< "charge_C" << " "
70-
<< "\n";
71-
} else {
72-
file_handler << "step" << " " << "s" << " "
73-
<< "x_mean" << " " << "x_min" << " " << "x_max" << " "
74-
<< "y_mean" << " " << "y_min" << " " << "y_max" << " "
75-
<< "t_mean" << " " << "t_min" << " " << "t_max" << " "
76-
<< "sig_x" << " " << "sig_y" << " " << "sig_t" << " "
77-
<< "px_mean" << " " << "px_min" << " " << "px_max" << " "
78-
<< "py_mean" << " " << "py_min" << " " << "py_max" << " "
79-
<< "pt_mean" << " " << "pt_min" << " " << "pt_max" << " "
80-
<< "sig_px" << " " << "sig_py" << " " << "sig_pt" << " "
81-
<< "emittance_x" << " " << "emittance_y" << " " << "emittance_t" << " "
82-
<< "alpha_x" << " " << "alpha_y" << " " << "alpha_t" << " "
83-
<< "beta_x" << " " << "beta_y" << " " << "beta_t" << " "
84-
<< "dispersion_x" << " " << "dispersion_px" << " "
85-
<< "dispersion_y" << " " << "dispersion_py" << " "
86-
<< "emittance_xn" << " " << "emittance_yn" << " " << "emittance_tn" << " "
87-
<< "charge_C" << " "
88-
<< "\n";
89-
}
90-
}
84+
if (!append)
85+
{
86+
write_column_header(file_handler, otype, has_charge);
9187
}
88+
}
89+
90+
void
91+
write_ref (
92+
amrex::AllPrintToFile & file_handler,
93+
impactx::RefPart const & ref_part,
94+
int step
95+
)
96+
{
97+
amrex::ParticleReal const s = ref_part.s;
98+
amrex::ParticleReal const beta = ref_part.beta();
99+
amrex::ParticleReal const gamma = ref_part.gamma();
100+
amrex::ParticleReal const beta_gamma = ref_part.beta_gamma();
101+
amrex::ParticleReal const x = ref_part.x;
102+
amrex::ParticleReal const y = ref_part.y;
103+
amrex::ParticleReal const z = ref_part.z;
104+
amrex::ParticleReal const t = ref_part.t;
105+
amrex::ParticleReal const px = ref_part.px;
106+
amrex::ParticleReal const py = ref_part.py;
107+
amrex::ParticleReal const pz = ref_part.pz;
108+
amrex::ParticleReal const pt = ref_part.pt;
109+
110+
// write particle data to file
111+
file_handler
112+
<< step << " " << s << " "
113+
<< beta << " " << gamma << " " << beta_gamma << " "
114+
<< x << " " << y << " " << z << " " << t << " "
115+
<< px << " " << py << " " << pz << " " << pt << "\n";
116+
}
117+
118+
void
119+
write_rbc (
120+
amrex::AllPrintToFile & file_handler,
121+
std::unordered_map<std::string, amrex::ParticleReal> const & rbc,
122+
amrex::ParticleReal s,
123+
int step
124+
)
125+
{
126+
// determine whether to output eigenemittances
127+
amrex::ParmParse pp_diag("diag");
128+
bool compute_eigenemittances = false;
129+
pp_diag.queryAdd("eigenemittances", compute_eigenemittances);
130+
131+
file_handler << step << " " << s << " "
132+
<< rbc.at("sig_x") << " " << rbc.at("sig_y") << " " << rbc.at("sig_t") << " "
133+
<< rbc.at("sig_px") << " " << rbc.at("sig_py") << " " << rbc.at("sig_pt") << " "
134+
<< rbc.at("emittance_x") << " " << rbc.at("emittance_y") << " " << rbc.at("emittance_t") << " "
135+
<< rbc.at("alpha_x") << " " << rbc.at("alpha_y") << " " << rbc.at("alpha_t") << " "
136+
<< rbc.at("beta_x") << " " << rbc.at("beta_y") << " " << rbc.at("beta_t") << " "
137+
<< rbc.at("dispersion_x") << " " << rbc.at("dispersion_px") << " "
138+
<< rbc.at("dispersion_y") << " " << rbc.at("dispersion_py") << " "
139+
<< rbc.at("emittance_xn") << " " << rbc.at("emittance_yn") << " " << rbc.at("emittance_tn");
140+
if (compute_eigenemittances) {
141+
file_handler << " "
142+
<< rbc.at("emittance_1") << " " << rbc.at("emittance_2") << " " << rbc.at("emittance_3");
143+
144+
}
145+
file_handler << "\n";
146+
}
147+
}
148+
149+
150+
namespace impactx::diagnostics
151+
{
152+
void DiagnosticOutput (
153+
ImpactXParticleContainer const & pc,
154+
OutputType const otype,
155+
std::string file_name,
156+
int step,
157+
bool append
158+
)
159+
{
160+
BL_PROFILE("impactx::diagnostics::DiagnosticOutput(pc)");
161+
162+
using namespace amrex::literals; // for _rt and _prt
92163

93-
if (otype == OutputType::PrintRefParticle) {
94-
// preparing to access reference particle data: RefPart
164+
// keep file open as we add more and more lines
165+
amrex::AllPrintToFile file_handler(std::move(file_name));
166+
prepare_header(file_handler, otype, true, append);
167+
168+
if (otype == OutputType::PrintRefParticle)
169+
{
95170
RefPart const ref_part = pc.GetRefParticle();
171+
write_ref(file_handler, ref_part, step);
172+
}
96173

97-
amrex::ParticleReal const s = ref_part.s;
98-
amrex::ParticleReal const beta = ref_part.beta();
99-
amrex::ParticleReal const gamma = ref_part.gamma();
100-
amrex::ParticleReal const beta_gamma = ref_part.beta_gamma();
101-
amrex::ParticleReal const x = ref_part.x;
102-
amrex::ParticleReal const y = ref_part.y;
103-
amrex::ParticleReal const z = ref_part.z;
104-
amrex::ParticleReal const t = ref_part.t;
105-
amrex::ParticleReal const px = ref_part.px;
106-
amrex::ParticleReal const py = ref_part.py;
107-
amrex::ParticleReal const pz = ref_part.pz;
108-
amrex::ParticleReal const pt = ref_part.pt;
109-
110-
// write particle data to file
111-
file_handler
112-
<< step << " " << s << " "
113-
<< beta << " " << gamma << " " << beta_gamma << " "
114-
<< x << " " << y << " " << z << " " << t << " "
115-
<< px << " " << py << " " << pz << " " << pt << "\n";
116-
} // if( otype == OutputType::PrintRefParticle)
117-
else if (otype == OutputType::PrintReducedBeamCharacteristics) {
174+
else if (otype == OutputType::PrintReducedBeamCharacteristics)
175+
{
176+
amrex::ParticleReal const s = pc.GetRefParticle().s;
118177
std::unordered_map<std::string, amrex::ParticleReal> const rbc =
119178
diagnostics::reduced_beam_characteristics(pc);
120179

121-
amrex::ParticleReal const s = pc.GetRefParticle().s;
122-
123-
// determine whether to output eigenemittances
124-
amrex::ParmParse pp_diag("diag");
125-
bool compute_eigenemittances = false;
126-
pp_diag.queryAdd("eigenemittances", compute_eigenemittances);
127-
128-
if (compute_eigenemittances) {
129-
file_handler << step << " " << s << " "
130-
<< rbc.at("x_mean") << " " << rbc.at("x_min") << " " << rbc.at("x_max") << " "
131-
<< rbc.at("y_mean") << " " << rbc.at("y_min") << " " << rbc.at("y_max") << " "
132-
<< rbc.at("t_mean") << " " << rbc.at("t_min") << " " << rbc.at("t_max") << " "
133-
<< rbc.at("sig_x") << " " << rbc.at("sig_y") << " " << rbc.at("sig_t") << " "
134-
<< rbc.at("px_mean") << " " << rbc.at("px_min") << " " << rbc.at("px_max") << " "
135-
<< rbc.at("py_mean") << " " << rbc.at("py_min") << " " << rbc.at("py_max") << " "
136-
<< rbc.at("pt_mean") << " " << rbc.at("pt_min") << " " << rbc.at("pt_max") << " "
137-
<< rbc.at("sig_px") << " " << rbc.at("sig_py") << " " << rbc.at("sig_pt") << " "
138-
<< rbc.at("emittance_x") << " " << rbc.at("emittance_y") << " " << rbc.at("emittance_t") << " "
139-
<< rbc.at("alpha_x") << " " << rbc.at("alpha_y") << " " << rbc.at("alpha_t") << " "
140-
<< rbc.at("beta_x") << " " << rbc.at("beta_y") << " " << rbc.at("beta_t") << " "
141-
<< rbc.at("dispersion_x") << " " << rbc.at("dispersion_px") << " "
142-
<< rbc.at("dispersion_y") << " " << rbc.at("dispersion_py") << " "
143-
<< rbc.at("emittance_xn") << " " << rbc.at("emittance_yn") << " " << rbc.at("emittance_tn") << " "
144-
<< rbc.at("emittance_1") << " " << rbc.at("emittance_2") << " " << rbc.at("emittance_3") << " "
145-
<< rbc.at("charge_C") << "\n";
146-
} else {
147-
file_handler << step << " " << s << " "
148-
<< rbc.at("x_mean") << " " << rbc.at("x_min") << " " << rbc.at("x_max") << " "
149-
<< rbc.at("y_mean") << " " << rbc.at("y_min") << " " << rbc.at("y_max") << " "
150-
<< rbc.at("t_mean") << " " << rbc.at("t_min") << " " << rbc.at("t_max") << " "
151-
<< rbc.at("sig_x") << " " << rbc.at("sig_y") << " " << rbc.at("sig_t") << " "
152-
<< rbc.at("px_mean") << " " << rbc.at("px_min") << " " << rbc.at("px_max") << " "
153-
<< rbc.at("py_mean") << " " << rbc.at("py_min") << " " << rbc.at("py_max") << " "
154-
<< rbc.at("pt_mean") << " " << rbc.at("pt_min") << " " << rbc.at("pt_max") << " "
155-
<< rbc.at("sig_px") << " " << rbc.at("sig_py") << " " << rbc.at("sig_pt") << " "
156-
<< rbc.at("emittance_x") << " " << rbc.at("emittance_y") << " " << rbc.at("emittance_t") << " "
157-
<< rbc.at("alpha_x") << " " << rbc.at("alpha_y") << " " << rbc.at("alpha_t") << " "
158-
<< rbc.at("beta_x") << " " << rbc.at("beta_y") << " " << rbc.at("beta_t") << " "
159-
<< rbc.at("dispersion_x") << " " << rbc.at("dispersion_px") << " "
160-
<< rbc.at("dispersion_y") << " " << rbc.at("dispersion_py") << " "
161-
<< rbc.at("emittance_xn") << " " << rbc.at("emittance_yn") << " " << rbc.at("emittance_tn") << " "
162-
<< rbc.at("charge_C") << "\n";
163-
}
164-
} // if( otype == OutputType::PrintReducedBeamCharacteristics)
180+
write_rbc(file_handler, rbc, s, step);
181+
}
165182

166-
// TODO: add as an option to the monitor element
167-
if (otype == OutputType::PrintNonlinearLensInvariants) {
183+
if (otype == OutputType::PrintNonlinearLensInvariants)
184+
{
168185
// create a host-side particle buffer
169186
auto tmp = pc.make_alike<amrex::PinnedArenaAllocator>();
170187

@@ -231,4 +248,39 @@ namespace impactx::diagnostics
231248
}
232249
}
233250

251+
void DiagnosticOutput (
252+
Map6x6 const & cm,
253+
RefPart const & ref_part,
254+
OutputType const otype,
255+
std::string file_name,
256+
int step,
257+
bool append
258+
)
259+
{
260+
BL_PROFILE("impactx::diagnostics::DiagnosticOutput(cm)");
261+
262+
// keep file open as we add more and more lines
263+
amrex::AllPrintToFile file_handler(std::move(file_name));
264+
prepare_header(file_handler, otype, false, append);
265+
266+
if (otype == OutputType::PrintRefParticle)
267+
{
268+
write_ref(file_handler, ref_part, step);
269+
}
270+
271+
else if (otype == OutputType::PrintReducedBeamCharacteristics)
272+
{
273+
amrex::ParticleReal const s = ref_part.s;
274+
std::unordered_map<std::string, amrex::ParticleReal> const rbc =
275+
diagnostics::reduced_beam_characteristics(cm, ref_part);
276+
277+
write_rbc(file_handler, rbc, s, step);
278+
}
279+
280+
if (otype == OutputType::PrintNonlinearLensInvariants)
281+
{
282+
throw std::runtime_error("DiagnosticOutput(cm): PrintNonlinearLensInvariants is not implemented.");
283+
}
284+
}
285+
234286
} // namespace impactx::diagnostics

0 commit comments

Comments
 (0)