forked from BLAST-WarpX/warpx
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathWarpXOpenPMD.cpp
1552 lines (1400 loc) · 60.9 KB
/
WarpXOpenPMD.cpp
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
/* Copyright 2019-2021 Axel Huebl, Junmin Gu
*
* This file is part of WarpX.
*
* License: BSD-3-Clause-LBNL
*/
#include "WarpXOpenPMD.H"
#include "Particles/ParticleIO.H"
#include "Diagnostics/ParticleDiag/ParticleDiag.H"
#include "FieldIO.H"
#include "Particles/Filter/FilterFunctors.H"
#include "Utils/TextMsg.H"
#include "Utils/Parser/ParserUtils.H"
#include "Utils/RelativeCellPosition.H"
#include "Utils/WarpXAlgorithmSelection.H"
#include "Utils/WarpXProfilerWrapper.H"
#include "WarpX.H"
#include "OpenPMDHelpFunction.H"
#include <ablastr/particles/IndexHandling.H>
#include <ablastr/warn_manager/WarnManager.H>
#include <AMReX.H>
#include <AMReX_ArrayOfStructs.H>
#include <AMReX_BLassert.H>
#include <AMReX_Box.H>
#include <AMReX_Config.H>
#include <AMReX_DataAllocator.H>
#include <AMReX_FArrayBox.H>
#include <AMReX_FabArray.H>
#include <AMReX_GpuQualifiers.H>
#include <AMReX_IntVect.H>
#include <AMReX_MFIter.H>
#include <AMReX_MultiFab.H>
#include <AMReX_PODVector.H>
#include <AMReX_ParallelDescriptor.H>
#include <AMReX_ParallelReduce.H>
#include <AMReX_Particle.H>
#include <AMReX_Particles.H>
#include <AMReX_Periodicity.H>
#include <AMReX_StructOfArrays.H>
#include <algorithm>
#include <cctype>
#include <cstdint>
#include <iostream>
#include <map>
#include <memory>
#include <regex>
#include <set>
#include <string>
#include <tuple>
#include <utility>
namespace detail
{
#ifdef WARPX_USE_OPENPMD
/** \brief Convert a snake_case string to a camelCase one.
*
* WarpX uses snake_case internally for some component
* names, but OpenPMD assumes "_" indicates vector or
* tensor fields.
*
* @return camelCase version of input
*/
inline std::string
snakeToCamel (const std::string& snake_string)
{
std::string camelString = snake_string;
const int n = camelString.length();
for (int x = 0; x < n; x++)
{
if (x == 0)
{
std::transform(camelString.begin(), camelString.begin()+1, camelString.begin(),
[](unsigned char c){ return std::tolower(c); });
}
if (camelString[x] == '_')
{
std::string tempString = camelString.substr(x + 1, 1);
std::transform(tempString.begin(), tempString.end(), tempString.begin(),
[](unsigned char c){ return std::toupper(c); });
camelString.erase(x, 2);
camelString.insert(x, tempString);
}
}
return camelString;
}
/** Create the option string
*
* @return JSON option string for openPMD::Series
*/
inline std::string
getSeriesOptions (std::string const & operator_type,
std::map< std::string, std::string > const & operator_parameters,
std::string const & engine_type,
std::map< std::string, std::string > const & engine_parameters)
{
if (operator_type.empty() && engine_type.empty())
return "{}";
std::string options;
std::string top_block;
std::string end_block;
std::string op_block;
std::string en_block;
std::string op_parameters;
for (const auto& kv : operator_parameters) {
if (!op_parameters.empty()) op_parameters.append(",\n");
op_parameters.append(std::string(12, ' ')) /* just pretty alignment */
.append("\"").append(kv.first).append("\": ") /* key */
.append("\"").append(kv.second).append("\""); /* value (as string) */
}
std::string en_parameters;
for (const auto& kv : engine_parameters) {
if (!en_parameters.empty()) en_parameters.append(",\n");
en_parameters.append(std::string(12, ' ')) /* just pretty alignment */
.append("\"").append(kv.first).append("\": ") /* key */
.append("\"").append(kv.second).append("\""); /* value (as string) */
}
// create the outer-level blocks
top_block = R"END(
{
"adios2": {)END";
end_block = R"END(
}
})END";
// add the operator string block
if (!operator_type.empty()) {
op_block = R"END(
"dataset": {
"operators": [
{
"type": ")END";
op_block += operator_type + "\"";
if (!op_parameters.empty()) {
op_block += R"END(,
"parameters": {
)END";
op_block += op_parameters +
"\n }";
}
op_block += R"END(
}
]
})END";
if (!engine_type.empty() || !en_parameters.empty())
op_block += ",";
} // end operator string block
// add the engine string block
if (!engine_type.empty() || !en_parameters.empty())
{
en_block = R"END(
"engine": {)END";
// non-default engine type
if (!engine_type.empty()) {
en_block += R"END(
"type": ")END";
en_block += engine_type + "\"";
if(!en_parameters.empty())
en_block += ",";
}
// non-default engine parameters
if (!en_parameters.empty()) {
en_block += R"END(
"parameters": {
)END";
en_block += en_parameters +
"\n }";
}
en_block += R"END(
})END";
} // end engine string block
options = top_block + op_block + en_block + end_block;
return options;
}
/** Unclutter a real_names to openPMD record
*
* @param fullName name as in real_names variable
* @return pair of openPMD record and component name
*/
inline std::pair< std::string, std::string >
name2openPMD ( std::string const& fullName )
{
std::string record_name = fullName;
std::string component_name = openPMD::RecordComponent::SCALAR;
// we use "_" as separator in names to group vector records
const std::size_t startComp = fullName.find_last_of('_');
if( startComp != std::string::npos ) { // non-scalar
record_name = fullName.substr(0, startComp);
component_name = fullName.substr(startComp + 1u);
}
return make_pair(record_name, component_name);
}
/** Return the component labels for particle positions
*/
inline std::vector< std::string >
getParticlePositionComponentLabels (bool ignore_dims=false)
{
using vs = std::vector< std::string >;
vs positionComponents;
if (ignore_dims) {
positionComponents = vs{"x", "y", "z"};
} else {
#if defined(WARPX_DIM_1D_Z)
positionComponents = vs{"z"};
#elif defined(WARPX_DIM_XZ)
positionComponents = vs{"x", "z"};
#elif defined(WARPX_DIM_RZ)
// note: although we internally store particle positions
// for AMReX in r,z and a theta attribute, we
// actually need them for algorithms (e.g. push)
// and I/O in Cartesian.
// Other attributes like momentum are consequently
// stored in x,y,z internally.
positionComponents = vs{"x", "y", "z"};
#elif defined(WARPX_DIM_3D)
positionComponents = vs{"x", "y", "z"};
#else
# error Unknown WarpX dimensionality.
#endif
}
return positionComponents;
}
/** Return the axis (index) names of a mesh
*
* This will be returned in C order. This is inverse of the Fortran order
* of the index labels for the AMReX FArrayBox.
*
* @param var_in_theta_mode indicate if this field will be output with theta
* modes (instead of a reconstructed 2D slice)
*/
inline std::vector< std::string >
getFieldAxisLabels ([[maybe_unused]] bool const var_in_theta_mode)
{
using vs = std::vector< std::string >;
// Fortran order of the index labels for the AMReX FArrayBox
#if defined(WARPX_DIM_1D_Z)
vs const axisLabels{"z"}; // z varies fastest in memory
#elif defined(WARPX_DIM_XZ)
vs const axisLabels{"x", "z"}; // x varies fastest in memory
#elif defined(WARPX_DIM_RZ)
// when we write individual modes of a field (default)
vs const circAxisLabels{"r", "z"}; // r varies fastest in memory
// if we just write reconstructed 2D fields at theta=0
vs const cartAxisLabels{"x", "z"}; // x varies fastest in memory
vs const axisLabels = var_in_theta_mode ? circAxisLabels : cartAxisLabels;
#elif defined(WARPX_DIM_3D)
vs const axisLabels{"x", "y", "z"}; // x varies fastest in memory
#else
# error Unknown WarpX dimensionality.
#endif
// revert to C order (fastest varying index last)
return {axisLabels.rbegin(), axisLabels.rend()};
}
/** Return the component names of a mesh
*
* @param var_in_theta_mode indicate if this field will be output with theta
* modes (instead of a reconstructed 2D slice)
*/
inline std::vector< std::string >
getFieldComponentLabels (bool const var_in_theta_mode)
{
using vs = std::vector< std::string >;
if (var_in_theta_mode) {
// if we write individual modes
vs fieldComponents{"r", "t", "z"};
return fieldComponents;
} else {
// if we just write reconstructed fields at theta=0 or are Cartesian
// note: 1D3V and 2D3V simulations still have 3 components for the fields
vs fieldComponents{"x", "y", "z"};
return fieldComponents;
}
}
/** Get the openPMD physical dimensionality of a record
*
* @param record_name name of the openPMD record
* @return map with base quantities and power scaling
*/
inline std::map< openPMD::UnitDimension, double >
getUnitDimension ( std::string const & record_name )
{
if( (record_name == "position") || (record_name == "positionOffset") ) return {
{openPMD::UnitDimension::L, 1.}
};
else if( record_name == "momentum" ) return {
{openPMD::UnitDimension::L, 1.},
{openPMD::UnitDimension::M, 1.},
{openPMD::UnitDimension::T, -1.}
};
else if( record_name == "charge" ) return {
{openPMD::UnitDimension::T, 1.},
{openPMD::UnitDimension::I, 1.}
};
else if( record_name == "mass" ) return {
{openPMD::UnitDimension::M, 1.}
};
else if( record_name == "E" ) return {
{openPMD::UnitDimension::L, 1.},
{openPMD::UnitDimension::M, 1.},
{openPMD::UnitDimension::T, -3.},
{openPMD::UnitDimension::I, -1.},
};
else if( record_name == "B" ) return {
{openPMD::UnitDimension::M, 1.},
{openPMD::UnitDimension::I, -1.},
{openPMD::UnitDimension::T, -2.}
};
else return {};
}
/** \brief For a given field that is to be written to an openPMD file,
* set the metadata that indicates the physical unit.
*/
inline void
setOpenPMDUnit ( openPMD::Mesh mesh, const std::string field_name )
{
if (field_name[0] == 'E'){ // Electric field
mesh.setUnitDimension({
{openPMD::UnitDimension::L, 1},
{openPMD::UnitDimension::M, 1},
{openPMD::UnitDimension::T, -3},
{openPMD::UnitDimension::I, -1},
});
} else if (field_name[0] == 'B'){ // Magnetic field
mesh.setUnitDimension({
{openPMD::UnitDimension::M, 1},
{openPMD::UnitDimension::I, -1},
{openPMD::UnitDimension::T, -2}
});
} else if (field_name[0] == 'j'){ // current
mesh.setUnitDimension({
{openPMD::UnitDimension::L, -2},
{openPMD::UnitDimension::I, 1},
});
} else if (field_name.substr(0,3) == "rho"){ // charge density
mesh.setUnitDimension({
{openPMD::UnitDimension::L, -3},
{openPMD::UnitDimension::I, 1},
{openPMD::UnitDimension::T, 1},
});
}
}
#endif // WARPX_USE_OPENPMD
} // namespace detail
#ifdef WARPX_USE_OPENPMD
WarpXOpenPMDPlot::WarpXOpenPMDPlot (
openPMD::IterationEncoding ie,
std::string openPMDFileType,
std::string operator_type,
std::map< std::string, std::string > operator_parameters,
std::string engine_type,
std::map< std::string, std::string > engine_parameters,
std::vector<bool> fieldPMLdirections)
:m_Series(nullptr),
m_Encoding(ie),
m_OpenPMDFileType(std::move(openPMDFileType)),
m_fieldPMLdirections(std::move(fieldPMLdirections))
{
m_OpenPMDoptions = detail::getSeriesOptions(operator_type, operator_parameters,
engine_type, engine_parameters);
}
WarpXOpenPMDPlot::~WarpXOpenPMDPlot ()
{
if( m_Series )
{
m_Series->flush();
m_Series.reset( nullptr );
}
}
std::string
WarpXOpenPMDPlot::GetFileName (std::string& filepath)
{
filepath.append("/");
// transform paths for Windows
#ifdef _WIN32
filepath = openPMD::auxiliary::replace_all(filepath, "/", "\\");
#endif
std::string filename = "openpmd";
//
// OpenPMD supports timestepped names
//
if (m_Encoding == openPMD::IterationEncoding::fileBased) {
const std::string fileSuffix = std::string("_%0") + std::to_string(m_file_min_digits) + std::string("T");
filename = filename.append(fileSuffix);
}
filename.append(".").append(m_OpenPMDFileType);
filepath.append(filename);
return filename;
}
void WarpXOpenPMDPlot::SetStep (int ts, const std::string& dirPrefix, int file_min_digits,
bool isBTD)
{
WARPX_ALWAYS_ASSERT_WITH_MESSAGE(ts >= 0 , "openPMD iterations are unsigned");
m_dirPrefix = dirPrefix;
m_file_min_digits = file_min_digits;
if( ! isBTD ) {
if (m_CurrentStep >= ts) {
// note m_Series is reset in Init(), so using m_Series->iterations.contains(ts) is only able to check the
// last written step in m_Series's life time, but not other earlier written steps by other m_Series
ablastr::warn_manager::WMRecordWarning("Diagnostics",
" Warning from openPMD writer: Already written iteration:"
+ std::to_string(ts)
);
}
}
m_CurrentStep = ts;
Init(openPMD::Access::CREATE, isBTD);
}
void WarpXOpenPMDPlot::CloseStep (bool isBTD, bool isLastBTDFlush)
{
// default close is true
bool callClose = true;
// close BTD file only when isLastBTDFlush is true
if (isBTD and !isLastBTDFlush) callClose = false;
if (callClose) {
if (m_Series) {
GetIteration(m_CurrentStep, isBTD).close();
}
// create a little helper file for ParaView 5.9+
if (amrex::ParallelDescriptor::IOProcessor())
{
// see Init()
std::string filepath = m_dirPrefix;
std::string const filename = GetFileName(filepath);
std::ofstream pv_helper_file(m_dirPrefix + "/paraview.pmd");
pv_helper_file << filename << std::endl;
pv_helper_file.close();
}
}
}
void
WarpXOpenPMDPlot::Init (openPMD::Access access, bool isBTD)
{
if( isBTD && m_Series != nullptr )
return; // already open for this snapshot (aka timestep in lab frame)
// either for the next ts file,
// or init a single file for all ts
std::string filepath = m_dirPrefix;
GetFileName(filepath);
// close a previously open series before creating a new one
// see ADIOS1 limitation: https://github.com/openPMD/openPMD-api/pull/686
if ( m_Encoding == openPMD::IterationEncoding::fileBased )
m_Series = nullptr;
else if ( m_Series != nullptr )
return;
if (amrex::ParallelDescriptor::NProcs() > 1) {
#if defined(AMREX_USE_MPI)
m_Series = std::make_unique<openPMD::Series>(
filepath, access,
amrex::ParallelDescriptor::Communicator(),
m_OpenPMDoptions
);
m_MPISize = amrex::ParallelDescriptor::NProcs();
m_MPIRank = amrex::ParallelDescriptor::MyProc();
#else
WARPX_ABORT_WITH_MESSAGE("openPMD-api not built with MPI support!");
#endif
} else {
m_Series = std::make_unique<openPMD::Series>(filepath, access, m_OpenPMDoptions);
m_MPISize = 1;
m_MPIRank = 1;
}
m_Series->setIterationEncoding( m_Encoding );
// input file / simulation setup author
if( !WarpX::authors.empty())
m_Series->setAuthor( WarpX::authors );
// more natural naming for PIC
m_Series->setMeshesPath( "fields" );
// conform to ED-PIC extension of openPMD
uint32_t const openPMD_ED_PIC = 1u;
m_Series->setOpenPMDextension( openPMD_ED_PIC );
// meta info
m_Series->setSoftware( "WarpX", WarpX::Version() );
}
void
WarpXOpenPMDPlot::WriteOpenPMDParticles (const amrex::Vector<ParticleDiag>& particle_diags,
const amrex::Real time, const bool use_pinned_pc,
const bool isBTD, const bool isLastBTDFlush,
const amrex::Vector<int>& totalParticlesFlushedAlready)
{
WARPX_PROFILE("WarpXOpenPMDPlot::WriteOpenPMDParticles()");
for (unsigned i = 0, n = particle_diags.size(); i < n; ++i) {
WarpXParticleContainer* pc = particle_diags[i].getParticleContainer();
PinnedMemoryParticleContainer* pinned_pc = particle_diags[i].getPinnedParticleContainer();
if (isBTD || use_pinned_pc)
if (!pinned_pc->isDefined())
continue; // Skip to the next particle container
PinnedMemoryParticleContainer tmp = (isBTD || use_pinned_pc) ?
pinned_pc->make_alike<amrex::PinnedArenaAllocator>() :
pc->make_alike<amrex::PinnedArenaAllocator>();
// names of amrex::Real and int particle attributes in SoA data
amrex::Vector<std::string> real_names;
amrex::Vector<std::string> int_names;
amrex::Vector<int> int_flags;
amrex::Vector<int> real_flags;
// see openPMD ED-PIC extension for namings
// note: an underscore separates the record name from its component
// for non-scalar records
real_names.push_back("weighting");
real_names.push_back("momentum_x");
real_names.push_back("momentum_y");
real_names.push_back("momentum_z");
#ifdef WARPX_DIM_RZ
real_names.push_back("theta");
#endif
// get the names of the real comps
real_names.resize(tmp.NumRealComps());
auto runtime_rnames = tmp.getParticleRuntimeComps();
for (auto const& x : runtime_rnames)
{
real_names[x.second+PIdx::nattribs] = detail::snakeToCamel(x.first);
}
// plot any "extra" fields by default
real_flags = particle_diags[i].m_plot_flags;
real_flags.resize(tmp.NumRealComps(), 1);
// and the names
int_names.resize(tmp.NumIntComps());
auto runtime_inames = tmp.getParticleRuntimeiComps();
for (auto const& x : runtime_inames)
{
int_names[x.second+0] = detail::snakeToCamel(x.first);
}
// plot by default
int_flags.resize(tmp.NumIntComps(), 1);
const auto mass = pc->AmIA<PhysicalSpecies::photon>() ? PhysConst::m_e : pc->getMass();
RandomFilter const random_filter(particle_diags[i].m_do_random_filter,
particle_diags[i].m_random_fraction);
UniformFilter const uniform_filter(particle_diags[i].m_do_uniform_filter,
particle_diags[i].m_uniform_stride);
ParserFilter parser_filter(particle_diags[i].m_do_parser_filter,
utils::parser::compileParser<ParticleDiag::m_nvars>
(particle_diags[i].m_particle_filter_parser.get()),
pc->getMass(), time);
parser_filter.m_units = InputUnits::SI;
GeometryFilter const geometry_filter(particle_diags[i].m_do_geom_filter,
particle_diags[i].m_diag_domain);
if (isBTD || use_pinned_pc) {
tmp.copyParticles(*pinned_pc, true);
particlesConvertUnits(ConvertDirection::WarpX_to_SI, &tmp, mass);
} else {
particlesConvertUnits(ConvertDirection::WarpX_to_SI, pc, mass);
using SrcData = WarpXParticleContainer::ParticleTileType::ConstParticleTileDataType;
tmp.copyParticles(*pc,
[random_filter,uniform_filter,parser_filter,geometry_filter]
AMREX_GPU_HOST_DEVICE
(const SrcData& src, int ip, const amrex::RandomEngine& engine)
{
const SuperParticleType& p = src.getSuperParticle(ip);
return random_filter(p, engine) * uniform_filter(p, engine)
* parser_filter(p, engine) * geometry_filter(p, engine);
}, true);
particlesConvertUnits(ConvertDirection::SI_to_WarpX, pc, mass);
}
// real_names contains a list of all real particle attributes.
// real_flags is 1 or 0, whether quantity is dumped or not.
{
if (isBTD) {
DumpToFile(&tmp,
particle_diags[i].getSpeciesName(),
m_CurrentStep,
real_flags,
int_flags,
real_names, int_names,
pc->getCharge(), pc->getMass(),
isBTD, isLastBTDFlush,
totalParticlesFlushedAlready[i]
);
} else {
DumpToFile(&tmp,
particle_diags[i].getSpeciesName(),
m_CurrentStep,
real_flags,
int_flags,
real_names, int_names,
pc->getCharge(), pc->getMass(),
isBTD, isLastBTDFlush,
0
);
}
}
}
}
void
WarpXOpenPMDPlot::DumpToFile (ParticleContainer* pc,
const std::string& name,
int iteration,
const amrex::Vector<int>& write_real_comp,
const amrex::Vector<int>& write_int_comp,
const amrex::Vector<std::string>& real_comp_names,
const amrex::Vector<std::string>& int_comp_names,
amrex::ParticleReal const charge,
amrex::ParticleReal const mass,
const bool isBTD,
const bool isLastBTDFlush,
int ParticleFlushOffset) {
WARPX_ALWAYS_ASSERT_WITH_MESSAGE(m_Series != nullptr, "openPMD: series must be initialized");
AMREX_ALWAYS_ASSERT(write_real_comp.size() == pc->NumRealComps());
AMREX_ALWAYS_ASSERT(write_int_comp.size() == pc->NumIntComps());
AMREX_ALWAYS_ASSERT(real_comp_names.size() == pc->NumRealComps());
AMREX_ALWAYS_ASSERT(int_comp_names.size() == pc->NumIntComps());
WarpXParticleCounter counter(pc);
auto const num_dump_particles = counter.GetTotalNumParticles();
openPMD::Iteration currIteration = GetIteration(iteration, isBTD);
openPMD::ParticleSpecies currSpecies = currIteration.particles[name];
// prepare data structures the first time BTD has non-zero particles
// we set some of them to zero extent, so we need to time that well
bool const is_first_flush_with_particles = num_dump_particles > 0 && ParticleFlushOffset == 0;
// BTD: we flush multiple times to the same lab step and thus need to resize
// our declared particle output sizes
bool const is_resizing_flush = num_dump_particles > 0 && ParticleFlushOffset > 0;
// write structure & declare particles in this (lab) step empty:
// if not BTD, then this is the only (and last) time we flush to this step
// if BTD, then we may do this multiple times until it is the last BTD flush
bool const is_last_flush_to_step = !isBTD || (isBTD && isLastBTDFlush);
// well, even in BTD we have to recognize that some lab stations may have no
// particles - so we mark them empty at the end of station reconstruction
bool const is_last_flush_and_never_particles =
is_last_flush_to_step && num_dump_particles == 0 && ParticleFlushOffset == 0;
//
// prepare structure and meta-data
//
// define positions & offset structure
const unsigned long long NewParticleVectorSize = num_dump_particles + ParticleFlushOffset;
// we will set up empty particles unless it's BTD, where we might add some in a following buffer dump
// during this setup, we mark some particle properties as constant and potentially zero-sized
bool doParticleSetup = true;
if (isBTD)
doParticleSetup = is_first_flush_with_particles || is_last_flush_and_never_particles;
// this setup stage also implicitly calls "makeEmpty" if needed (i.e., is_last_flush_and_never_particles)
// for BTD, we call this multiple times as we may resize in subsequent dumps if number of particles in the buffer > 0
if (doParticleSetup || is_resizing_flush) {
SetupPos(currSpecies, NewParticleVectorSize, isBTD);
SetupRealProperties(pc, currSpecies, write_real_comp, real_comp_names, write_int_comp, int_comp_names,
NewParticleVectorSize, isBTD);
}
if (is_last_flush_to_step) {
SetConstParticleRecordsEDPIC(currSpecies, NewParticleVectorSize, charge, mass);
}
// open files from all processors, in case some will not contribute below
m_Series->flush();
// dump individual particles
bool contributed_particles = false; // did the local MPI rank contribute particles?
for (auto currentLevel = 0; currentLevel <= pc->finestLevel(); currentLevel++) {
uint64_t offset = static_cast<uint64_t>( counter.m_ParticleOffsetAtRank[currentLevel] );
// For BTD, the offset include the number of particles already flushed
if (isBTD) offset += ParticleFlushOffset;
for (ParticleIter pti(*pc, currentLevel); pti.isValid(); ++pti) {
auto const numParticleOnTile = pti.numParticles();
uint64_t const numParticleOnTile64 = static_cast<uint64_t>( numParticleOnTile );
// Do not call storeChunk() with zero-sized particle tiles:
// https://github.com/openPMD/openPMD-api/issues/1147
// https://github.com/ECP-WarpX/WarpX/pull/1898#discussion_r745008290
if (numParticleOnTile == 0) continue;
contributed_particles = true;
// get position and particle ID from aos
// note: this implementation iterates the AoS 4x...
// if we flush late as we do now, we can also copy out the data in one go
const auto &aos = pti.GetArrayOfStructs(); // size = numParticlesOnTile
{
// Save positions
#if defined(WARPX_DIM_RZ)
{
const std::shared_ptr<amrex::ParticleReal> z(
new amrex::ParticleReal[numParticleOnTile],
[](amrex::ParticleReal const *p) { delete[] p; }
);
for (auto i = 0; i < numParticleOnTile; i++)
z.get()[i] = aos[i].pos(1); // {0: "r", 1: "z"}
std::string const positionComponent = "z";
currSpecies["position"]["z"].storeChunk(z, {offset}, {numParticleOnTile64});
}
// reconstruct x and y from polar coordinates r, theta
auto const& soa = pti.GetStructOfArrays();
amrex::ParticleReal const* theta = soa.GetRealData(PIdx::theta).dataPtr();
WARPX_ALWAYS_ASSERT_WITH_MESSAGE(theta != nullptr, "openPMD: invalid theta pointer.");
WARPX_ALWAYS_ASSERT_WITH_MESSAGE(int(soa.GetRealData(PIdx::theta).size()) == numParticleOnTile,
"openPMD: theta and tile size do not match");
{
const std::shared_ptr< amrex::ParticleReal > x(
new amrex::ParticleReal[numParticleOnTile],
[](amrex::ParticleReal const *p){ delete[] p; }
);
const std::shared_ptr< amrex::ParticleReal > y(
new amrex::ParticleReal[numParticleOnTile],
[](amrex::ParticleReal const *p){ delete[] p; }
);
for (auto i=0; i<numParticleOnTile; i++) {
auto const r = aos[i].pos(0); // {0: "r", 1: "z"}
x.get()[i] = r * std::cos(theta[i]);
y.get()[i] = r * std::sin(theta[i]);
}
currSpecies["position"]["x"].storeChunk(x, {offset}, {numParticleOnTile64});
currSpecies["position"]["y"].storeChunk(y, {offset}, {numParticleOnTile64});
}
#else
auto const positionComponents = detail::getParticlePositionComponentLabels();
for (auto currDim = 0; currDim < AMREX_SPACEDIM; currDim++) {
const std::shared_ptr<amrex::ParticleReal> curr(
new amrex::ParticleReal[numParticleOnTile],
[](amrex::ParticleReal const *p) { delete[] p; }
);
for (auto i = 0; i < numParticleOnTile; i++) {
curr.get()[i] = aos[i].pos(currDim);
}
std::string const positionComponent = positionComponents[currDim];
currSpecies["position"][positionComponent].storeChunk(curr, {offset},
{numParticleOnTile64});
}
#endif
// save particle ID after converting it to a globally unique ID
const std::shared_ptr<uint64_t> ids(
new uint64_t[numParticleOnTile],
[](uint64_t const *p) { delete[] p; }
);
for (auto i = 0; i < numParticleOnTile; i++) {
ids.get()[i] = ablastr::particles::localIDtoGlobal(aos[i].id(), aos[i].cpu());
}
auto const scalar = openPMD::RecordComponent::SCALAR;
currSpecies["id"][scalar].storeChunk(ids, {offset}, {numParticleOnTile64});
}
// save "extra" particle properties in AoS and SoA
SaveRealProperty(pti,
currSpecies,
offset,
write_real_comp, real_comp_names,
write_int_comp, int_comp_names);
offset += numParticleOnTile64;
} // pti
} // currentLevel
// work-around for BTD particle resize in ADIOS2
//
// This issues an empty ADIOS2 Put to make sure the new global shape
// meta-data is committed for each variable.
//
// Refs.:
// https://github.com/ECP-WarpX/WarpX/issues/3389
// https://github.com/ornladios/ADIOS2/issues/3455
// BP4 (ADIOS 2.8): last MPI rank's `Put` meta-data wins
// BP5 (ADIOS 2.8): everyone has to write an empty block
if (is_resizing_flush && !contributed_particles && isBTD && m_Series->backend() == "ADIOS2") {
for( auto & [record_name, record] : currSpecies ) {
for( auto & [comp_name, comp] : record ) {
if (comp.constant()) continue;
auto dtype = comp.getDatatype();
switch (dtype) {
case openPMD::Datatype::FLOAT :
[[fallthrough]];
case openPMD::Datatype::DOUBLE : {
auto empty_data = std::make_shared<amrex::ParticleReal>();
comp.storeChunk(empty_data, {uint64_t(0)}, {uint64_t(0)});
break;
}
case openPMD::Datatype::UINT : {
auto empty_data = std::make_shared<unsigned int>();
comp.storeChunk(empty_data, {uint64_t(0)}, {uint64_t(0)});
break;
}
case openPMD::Datatype::ULONG : {
auto empty_data = std::make_shared<unsigned long>();
comp.storeChunk(empty_data, {uint64_t(0)}, {uint64_t(0)});
break;
}
case openPMD::Datatype::ULONGLONG : {
auto empty_data = std::make_shared<unsigned long long>();
comp.storeChunk(empty_data, {uint64_t(0)}, {uint64_t(0)});
break;
}
default : {
std::string msg = "WarpX openPMD ADIOS2 work-around has unknown dtype: ";
msg += datatypeToString(dtype);
WARPX_ABORT_WITH_MESSAGE(msg);
break;
}
}
}
}
}
m_Series->flush();
}
void
WarpXOpenPMDPlot::SetupRealProperties (ParticleContainer const * pc,
openPMD::ParticleSpecies& currSpecies,
const amrex::Vector<int>& write_real_comp,
const amrex::Vector<std::string>& real_comp_names,
const amrex::Vector<int>& write_int_comp,
const amrex::Vector<std::string>& int_comp_names,
const unsigned long long np, bool const isBTD) const
{
std::string options = "{}";
if (isBTD) options = "{ \"resizable\": true }";
auto dtype_real = openPMD::Dataset(openPMD::determineDatatype<amrex::ParticleReal>(), {np}, options);
auto dtype_int = openPMD::Dataset(openPMD::determineDatatype<int>(), {np}, options);
//
// the beam/input3d showed write_real_comp.size() = 16 while only 10 real comp names
// so using the min to be safe.
//
auto const getComponentRecord = [&currSpecies](std::string const comp_name) {
// handle scalar and non-scalar records by name
const auto [record_name, component_name] = detail::name2openPMD(comp_name);
return currSpecies[record_name][component_name];
};
auto const real_counter = std::min(write_real_comp.size(), real_comp_names.size());
for (int i = 0; i < real_counter; ++i) {
if (write_real_comp[i]) {
getComponentRecord(real_comp_names[i]).resetDataset(dtype_real);
}
}
auto const int_counter = std::min(write_int_comp.size(), int_comp_names.size());
for (int i = 0; i < int_counter; ++i) {
if (write_int_comp[i]) {
getComponentRecord(int_comp_names[i]).resetDataset(dtype_int);
}
}
std::set< std::string > addedRecords; // add meta-data per record only once
for (auto idx=0; idx<pc->NumRealComps(); idx++) {
auto ii = ParticleContainer::NStructReal + idx; // jump over extra AoS names
if (write_real_comp[ii]) {
// handle scalar and non-scalar records by name
const auto [record_name, component_name] = detail::name2openPMD(real_comp_names[ii]);
auto currRecord = currSpecies[record_name];
// meta data for ED-PIC extension
[[maybe_unused]] const auto [_, newRecord] = addedRecords.insert(record_name);
if( newRecord ) {
currRecord.setUnitDimension( detail::getUnitDimension(record_name) );
if( record_name == "weighting" )
currRecord.setAttribute( "macroWeighted", 1u );
else
currRecord.setAttribute( "macroWeighted", 0u );
if( record_name == "momentum" || record_name == "weighting" )
currRecord.setAttribute( "weightingPower", 1.0 );
else
currRecord.setAttribute( "weightingPower", 0.0 );
}
}
}
for (auto idx=0; idx<int_counter; idx++) {
auto ii = ParticleContainer::NStructInt + idx; // jump over extra AoS names
if (write_int_comp[ii]) {
// handle scalar and non-scalar records by name
const auto [record_name, component_name] = detail::name2openPMD(int_comp_names[ii]);
auto currRecord = currSpecies[record_name];
// meta data for ED-PIC extension
[[maybe_unused]] const auto [_, newRecord] = addedRecords.insert(record_name);
if( newRecord ) {
currRecord.setUnitDimension( detail::getUnitDimension(record_name) );
currRecord.setAttribute( "macroWeighted", 0u );
if( record_name == "momentum" || record_name == "weighting" )
currRecord.setAttribute( "weightingPower", 1.0 );
else
currRecord.setAttribute( "weightingPower", 0.0 );
}
}
}
}
void
WarpXOpenPMDPlot::SaveRealProperty (ParticleIter& pti,
openPMD::ParticleSpecies& currSpecies,
unsigned long long const offset,
amrex::Vector<int> const& write_real_comp,
amrex::Vector<std::string> const& real_comp_names,
amrex::Vector<int> const& write_int_comp,
amrex::Vector<std::string> const& int_comp_names) const
{
auto const numParticleOnTile = pti.numParticles();
uint64_t const numParticleOnTile64 = static_cast<uint64_t>( numParticleOnTile );
auto const& aos = pti.GetArrayOfStructs(); // size = numParticlesOnTile
auto const& soa = pti.GetStructOfArrays();
// first we concatinate the AoS into contiguous arrays
{
// note: WarpX does not yet use extra AoS Real attributes
for( auto idx=0; idx<ParticleIter::ContainerType::NStructReal; idx++ ) { // lgtm [cpp/constant-comparison]
if( write_real_comp[idx] ) {
// handle scalar and non-scalar records by name
const auto [record_name, component_name] = detail::name2openPMD(real_comp_names[idx]);
auto currRecord = currSpecies[record_name];
auto currRecordComp = currRecord[component_name];
const std::shared_ptr< amrex::ParticleReal > d(
new amrex::ParticleReal[numParticleOnTile],
[](amrex::ParticleReal const *p){ delete[] p; }
);
for( auto kk=0; kk<numParticleOnTile; kk++ )
d.get()[kk] = aos[kk].rdata(idx);
currRecordComp.storeChunk(d,
{offset}, {numParticleOnTile64});
}
}
}
auto const getComponentRecord = [&currSpecies](std::string const comp_name) {
// handle scalar and non-scalar records by name
const auto [record_name, component_name] = detail::name2openPMD(comp_name);
return currSpecies[record_name][component_name];
};
// here we the save the SoA properties (real)
{
auto const real_counter = std::min(write_real_comp.size(), real_comp_names.size());
for (auto idx=0; idx<real_counter; idx++) {
auto ii = ParticleIter::ContainerType::NStructReal + idx; // jump over extra AoS names
if (write_real_comp[ii]) {
getComponentRecord(real_comp_names[ii]).storeChunkRaw(
soa.GetRealData(idx).data(), {offset}, {numParticleOnTile64});
}
}
}
// and now SoA int properties
{
auto const int_counter = std::min(write_int_comp.size(), int_comp_names.size());
for (auto idx=0; idx<int_counter; idx++) {
auto ii = ParticleIter::ContainerType::NStructInt + idx; // jump over extra AoS names
if (write_int_comp[ii]) {
getComponentRecord(int_comp_names[ii]).storeChunkRaw(
soa.GetIntData(idx).data(), {offset}, {numParticleOnTile64});
}
}
}