Allow saving field attributes (like polarization)

AngelFP · Mar 10, 2023 · e7f0011 · e7f0011
1 parent 4a1856b
commit e7f0011
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 7 deletions.
diff --git a/wake_t/diagnostics/openpmd_diag.py b/wake_t/diagnostics/openpmd_diag.py
@@ -303,6 +303,8 @@ def _write_fields(self, it, wf_data):
             # properly defined in the openPMD standard.
             fld.geometry = Geometry.thetaMode  # Geometry.cylindrical
             fld.set_attribute('fieldSmoothing', 'none')
+            for attr, val in wf_data[field]['attributes'].items():
+                fld.set_attribute(attr, val)
             fld.axis_labels = wf_data[field]['grid']['labels']
             fld.grid_spacing = wf_data[field]['grid']['spacing']
             global_offset = deepcopy(wf_data[field]['grid']['global_offset'])

diff --git a/wake_t/fields/rz_wakefield.py b/wake_t/fields/rz_wakefield.py
@@ -173,6 +173,7 @@ def _get_openpmd_diagnostics_data(self, global_time):
         fld_position = [0.5, 0.]
         fld_names = ['E', 'B', 'rho']
         fld_comps = [['r', 't', 'z'], ['r', 't', 'z'], None]
+        fld_attrs = [{}, {}, {}]
         fld_arrays = [
             [np.ascontiguousarray(self.e_r.T[2:-2, 2:-2]),
              np.ascontiguousarray(self.e_t.T[2:-2, 2:-2]),
@@ -185,6 +186,7 @@ def _get_openpmd_diagnostics_data(self, global_time):
         if self.laser is not None:
             fld_names += ['a_mod', 'a_phase']
             fld_comps += [None, None]
+            fld_attrs += [{'polarization': self.laser.polarization}, {}]
             fld_arrays += [
                 [np.ascontiguousarray(np.abs(self.laser.get_envelope().T))],
                 [np.ascontiguousarray(np.angle(self.laser.get_envelope().T))]
@@ -193,9 +195,10 @@ def _get_openpmd_diagnostics_data(self, global_time):
 
         # Generate dictionary for openPMD diagnostics.
         diag_data = generate_field_diag_dictionary(
-            fld_names, fld_comps, fld_arrays, fld_comp_pos, grid_labels,
-            grid_spacing, grid_global_offset, fld_solver, fld_solver_params,
-            fld_boundary, fld_boundary_params, part_boundary,
-            part_boundary_params, current_smoothing, charge_correction)
+            fld_names, fld_comps, fld_attrs, fld_arrays, fld_comp_pos,
+            grid_labels, grid_spacing, grid_global_offset, fld_solver,
+            fld_solver_params, fld_boundary, fld_boundary_params,
+            part_boundary, part_boundary_params, current_smoothing,
+            charge_correction)
 
         return diag_data
diff --git a/wake_t/utilities/other.py b/wake_t/utilities/other.py
@@ -16,7 +16,7 @@ def print_progress_bar(pre_string, step, total_steps):
 
 
 def generate_field_diag_dictionary(
-        fld_names, fld_comps, fld_arrays, fld_comp_pos, grid_labels,
+        fld_names, fld_comps, fld_attrs, fld_arrays, fld_comp_pos, grid_labels,
         grid_spacing, grid_global_offset, fld_solver, fld_solver_params,
         fld_boundary, fld_boundary_params, part_boundary,
         part_boundary_params, current_smoothing, charge_correction):
@@ -27,8 +27,8 @@ def generate_field_diag_dictionary(
     """
     diag_data = {}
     diag_data['fields'] = fld_names
-    fld_zip = zip(fld_names, fld_comps, fld_arrays, fld_comp_pos)
-    for fld, comps, arrays, pos in fld_zip:
+    fld_zip = zip(fld_names, fld_comps, fld_attrs, fld_arrays, fld_comp_pos)
+    for fld, comps, attrs, arrays, pos in fld_zip:
         diag_data[fld] = {}
         if comps is not None:
             diag_data[fld]['comps'] = {}
@@ -43,6 +43,7 @@ def generate_field_diag_dictionary(
         diag_data[fld]['grid']['spacing'] = grid_spacing
         diag_data[fld]['grid']['labels'] = grid_labels
         diag_data[fld]['grid']['global_offset'] = grid_global_offset
+        diag_data[fld]['attributes'] = attrs
     diag_data['field_solver'] = fld_solver
     diag_data['field_solver_params'] = fld_solver_params
     diag_data['field_boundary'] = fld_boundary