diff options
Diffstat (limited to 'Tools/read_raw_data.py')
| -rw-r--r-- | Tools/read_raw_data.py | 45 |
1 files changed, 22 insertions, 23 deletions
diff --git a/Tools/read_raw_data.py b/Tools/read_raw_data.py index e9745f08b..deb1f36a2 100644 --- a/Tools/read_raw_data.py +++ b/Tools/read_raw_data.py @@ -9,18 +9,18 @@ _component_names = ['Ex', 'Ey', 'Ez', 'Bx', 'By', 'Bz', 'jx', 'jy', 'jz', 'rho'] def read_data(plt_file): ''' - This function reads the raw (i.e. not averaged to cell centers) data + This function reads the raw (i.e. not averaged to cell centers) data from a WarpX plt file. The plt file must have been written with the - plot_raw_fields option turned on, so that it contains a raw_data + plot_raw_fields option turned on, so that it contains a raw_data sub-directory. This is only really useful for single-level data. Arguments: plt_file : An AMReX plt_file file. Must contain a raw_data directory. - + Returns: - A list of dictionaries where the keys are field name strings and the values + A list of dictionaries where the keys are field name strings and the values are numpy arrays. Each entry in the list corresponds to a different level. Example: @@ -34,7 +34,7 @@ def read_data(plt_file): raw_files = glob(plt_file + "/raw_fields/Level_*/") for raw_file in raw_files: field_names = _get_field_names(raw_file) - + data = {} for field in field_names: data[field] = _read_field(raw_file, field) @@ -48,19 +48,19 @@ def read_lab_snapshot(snapshot): ''' This reads the data from one of the lab frame snapshots generated when - WarpX is run with boosted frame diagnostics turned on. It returns a + WarpX is run with boosted frame diagnostics turned on. It returns a dictionary of numpy arrays, where each key corresponds to one of the data fields ("Ex", "By,", etc... ). These values are cell-centered. ''' - hdrs = glob(snapshot + "/Level_0/buffer?????_H") + hdrs = glob(snapshot + "/Level_0/buffer??????_H") hdrs.sort() boxes, file_names, offsets, header = _read_header(hdrs[0]) dom_lo, dom_hi = _combine_boxes(boxes) domain_size = dom_hi - dom_lo + 1 - + space_dim = len(dom_lo) if space_dim == 2: direction = 1 @@ -69,14 +69,14 @@ def read_lab_snapshot(snapshot): buffer_data = _read_buffer(snapshot, hdrs[0]) buffer_size = buffer_data['Bx'].shape[direction] - + data = {} for i in range(header.ncomp): if space_dim == 3: data[_component_names[i]] = np.zeros((domain_size[0], domain_size[1], buffer_size*len(hdrs))) elif space_dim == 2: data[_component_names[i]] = np.zeros((domain_size[0], buffer_size*len(hdrs))) - + for i, hdr in enumerate(hdrs): buffer_data = _read_buffer(snapshot, hdr) if data is None: @@ -87,7 +87,7 @@ def read_lab_snapshot(snapshot): data[k][:,:,buffer_size*i:buffer_size*(i+1)] = v[:,:,:] elif space_dim == 2: data[k][:,buffer_size*i:buffer_size*(i+1)] = v[:,:] - + return data @@ -103,7 +103,7 @@ def _string_to_numpy_array(s): def _line_to_numpy_arrays(line): lo_corner = _string_to_numpy_array(line[0][1:]) hi_corner = _string_to_numpy_array(line[1][:]) - node_type = _string_to_numpy_array(line[2][:-1]) + node_type = _string_to_numpy_array(line[2][:-1]) return lo_corner, hi_corner, node_type @@ -152,39 +152,39 @@ def _combine_boxes(boxes): def _read_field(raw_file, field_name): - header_file = raw_file + field + "_H" + header_file = raw_file + field_name + "_H" boxes, file_names, offsets, header = _read_header(header_file) ng = header.nghost - lo, hi = _combine_boxes(boxes) - data = np.zeros(hi - lo + 1) + dom_lo, dom_hi = _combine_boxes(boxes) + data = np.zeros(dom_hi - dom_lo + 1) for box, fn, offset in zip(boxes, file_names, offsets): - lo = box[0] - hi = box[1] + lo = box[0] - dom_lo + hi = box[1] - dom_lo shape = hi - lo + 1 with open(raw_file + fn, "rb") as f: f.seek(offset) f.readline() # always skip the first line arr = np.fromfile(f, 'float64', np.product(shape)) arr = arr.reshape(shape, order='F') - data[[slice(l,h+1) for l, h in zip(lo+ng, hi+ng)]] = arr + data[[slice(l,h+1) for l, h in zip(lo, hi)]] = arr return data def _read_buffer(snapshot, header_fn): - + boxes, file_names, offsets, header = _read_header(header_fn) ng = header.nghost dom_lo, dom_hi = _combine_boxes(boxes) - + all_data = {} for i in range(header.ncomp): all_data[_component_names[i]] = np.zeros(dom_hi - dom_lo + 1) - + for box, fn, offset in zip(boxes, file_names, offsets): lo = box[0] - dom_lo hi = box[1] - dom_lo @@ -213,7 +213,7 @@ if __name__ == "__main__": for level in range(2): for name, vals in data[level].items(): print(level, name, vals.shape, vals.min(), vals.max()) - + # make a projection along the z-axis of the 'jx' field for level 0 level = 0 plt.pcolormesh(data[level]['jx'].sum(axis=2)) @@ -223,4 +223,3 @@ if __name__ == "__main__": level = 1 plt.pcolormesh(data[level]['Bx_cp'].sum(axis=0)) plt.savefig('Bx_cp') - |