From 7418556773c74827478767d2a58c9220cf5900b3 Mon Sep 17 00:00:00 2001
From: arnav-singhal <43693924+arnav-singhal@users.noreply.github.com>
Date: Fri, 28 Jun 2024 11:22:13 -0700
Subject: [PATCH] HDF5: Correctly handle writing only some components when
 writing Particle (#4005)

## Summary
Correctly names file attribute to match component index with dataset
index and correctly calculates size of each particle's data so that the
dataset is aligned.
## Additional background
Previously, all components of a particle were being counted, whether
written or not, leading to misalignment across different processes when
writing.
## Checklist

The proposed changes:
- [x] fix a bug or incorrect behavior in AMReX
- [ ] add new capabilities to AMReX
- [ ] changes answers in the test suite to more than roundoff level
- [ ] are likely to significantly affect the results of downstream AMReX
users
- [ ] include documentation in the code and/or rst files, if appropriate
---
 Src/Extern/HDF5/AMReX_ParticleHDF5.H            | 17 +++++++++++++++--
 .../HDF5/AMReX_WriteBinaryParticleDataHDF5.H    | 11 +++++++++--
 2 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/Src/Extern/HDF5/AMReX_ParticleHDF5.H b/Src/Extern/HDF5/AMReX_ParticleHDF5.H
index 398b84f7607..00fdf0c14be 100644
--- a/Src/Extern/HDF5/AMReX_ParticleHDF5.H
+++ b/Src/Extern/HDF5/AMReX_ParticleHDF5.H
@@ -705,19 +705,32 @@ ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator, CellAssig
     hsize_t        total_mfi = 0, total_real_size = 0, total_int_size = 0, real_file_offset = 0, int_file_offset = 0;
     hsize_t        my_int_offset, my_int_count, my_real_offset, my_real_count;
 
+    // Count total number of components written
+    int real_comp_count = AMREX_SPACEDIM; // position values
+
+    for (int i = 0; i < NStructReal + NumRealComps(); ++i ) {
+        if (write_real_comp[i]) { ++real_comp_count; }
+    }
+
+    int int_comp_count = 2; // cpu and id values
+
+    for (int i = 0; i < NStructInt + NumIntComps(); ++i ) {
+        if (write_int_comp[i]) { ++int_comp_count; }
+    }
+
     // Get the size for each mf so we know the amount of data from each rank
     for (MFIter mfi(state); mfi.isValid(); ++mfi) {
         const int grid = mfi.index();
         if (count[grid] == 0)
             continue;
 
-        int_size  = count[grid] * (2 + NStructInt + NumIntComps());
+        int_size = count[grid] * int_comp_count;
         my_mfi_int_size.push_back(int_size);
         my_nparticles.push_back(count[grid]);
         my_mfi_int_total_size += int_size;
 
 
-        real_size = count[grid] * (AMREX_SPACEDIM + NStructReal + NumRealComps());
+        real_size = count[grid] * real_comp_count;
         my_mfi_real_size.push_back(real_size);
         my_mfi_real_total_size += real_size;
         my_mfi_cnt++;
diff --git a/Src/Extern/HDF5/AMReX_WriteBinaryParticleDataHDF5.H b/Src/Extern/HDF5/AMReX_WriteBinaryParticleDataHDF5.H
index eb13fc5abb1..271f2e0e4cd 100644
--- a/Src/Extern/HDF5/AMReX_WriteBinaryParticleDataHDF5.H
+++ b/Src/Extern/HDF5/AMReX_WriteBinaryParticleDataHDF5.H
@@ -289,12 +289,16 @@ void WriteHDF5ParticleDataSync (PC const& pc,
         CreateWriteHDF5Attr(fid, "num_component_real", 1, &num_output_real, H5T_NATIVE_INT);
 
         char comp_name[128];
+
         // Real component names
+        int real_comp_count = AMREX_SPACEDIM; // position values
+
         for (int i = 0; i < NStructReal + pc.NumRealComps(); ++i ) {
             if (write_real_comp[i]) {
                 /* HdrFile << real_comp_names[i] << '\n'; */
-                snprintf(comp_name, sizeof comp_name, "real_component_%d", i);
+                snprintf(comp_name, sizeof comp_name, "real_component_%d", real_comp_count);
                 CreateWriteHDF5AttrString(fid, comp_name, real_comp_names[i].c_str());
+                ++real_comp_count;
             }
         }
 
@@ -302,10 +306,13 @@ void WriteHDF5ParticleDataSync (PC const& pc,
         CreateWriteHDF5Attr(fid, "num_component_int", 1, &num_output_int, H5T_NATIVE_INT);
 
         // int component names
+        int int_comp_count = 2;
+
         for (int i = 0; i < NStructInt + pc.NumIntComps(); ++i ) {
             if (write_int_comp[i]) {
-                snprintf(comp_name, sizeof comp_name, "int_component_%d", i);
+                snprintf(comp_name, sizeof comp_name, "int_component_%d", int_comp_count);
                 CreateWriteHDF5AttrString(fid, comp_name, int_comp_names[i].c_str());
+                ++int_comp_count;
             }
         }