# Copyright 2017-2019 Andrew Myers, Axel Huebl, Maxence Thevenet # Remi Lehe # # This file is part of WarpX. # # License: BSD-3-Clause-LBNL from glob import glob import os import numpy as np from collections import namedtuple HeaderInfo = namedtuple('HeaderInfo', ['version', 'how', 'ncomp', 'nghost']) def read_data(plt_file): ''' 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 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 are numpy arrays. Each entry in the list corresponds to a different level. Example: >>> data = read_data("plt00016") >>> print(data.keys()) >>> print(data['Ex'].shape) ''' all_data = [] 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) all_data.append(data) return all_data def read_lab_snapshot(snapshot, global_header): ''' 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 dictionary of numpy arrays, where each key corresponds to one of the data fields ("Ex", "By,", etc... ). These values are cell-centered. ''' global_info = _read_global_Header(global_header) 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) local_info = _read_local_Header(snapshot + "/Header", space_dim) ncellz_snapshots = local_info['nz'] dzcell_snapshots = (local_info['zmax']-local_info['zmin'])/local_info['nz'] _component_names = local_info['field_names'] field1 = _component_names[0] if space_dim == 2: direction = 1 else: direction = 2 buffer_fullsize = 0 buffer_allsizes = [0] for i, hdr in enumerate(hdrs): buffer_data = _read_buffer(snapshot, hdr, _component_names) buffer_fullsize += buffer_data[field1].shape[direction] buffer_allsizes.append(buffer_data[field1].shape[direction]) buffer_allstarts = np.cumsum(buffer_allsizes) data = {} for i in range(header.ncomp): if space_dim == 3: data[_component_names[i]] = np.zeros((domain_size[0], domain_size[1], buffer_fullsize)) elif space_dim == 2: data[_component_names[i]] = np.zeros((domain_size[0], buffer_fullsize)) for i, hdr in enumerate(hdrs): buffer_data = _read_buffer(snapshot, hdr, _component_names) if data is None: data = buffer_data else: for k,v in buffer_data.items(): data[k][..., buffer_allstarts[i]:buffer_allstarts[i+1]] = v[...] info = local_info # Add some handy info x = np.linspace(local_info['xmin'], local_info['xmax'], local_info['nx']) y = np.linspace(local_info['ymin'], local_info['ymax'], local_info['ny']) z = np.linspace(local_info['zmin'], local_info['zmax'], local_info['nz']) info.update({ 'x' : x, 'y' : y, 'z' : z }) return data, info # For the moment, the back-transformed diagnostics must be read with # custom functions like this one. # It should be OpenPMD-compliant hdf5 files soon, making this part outdated. def get_particle_field(snapshot, species, field): fn = snapshot + '/' + species files = glob(os.path.join(fn, field + '_*')) files.sort() all_data = np.array([]) for f in files: data = np.fromfile(f) all_data = np.concatenate((all_data, data)) return all_data def _get_field_names(raw_file): header_files = glob(raw_file + "*_H") return [hf.split("/")[-1][:-2] for hf in header_files] def _string_to_numpy_array(s): return np.array([int(v) for v in s[1:-1].split(",")], dtype=np.int64) 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]) return lo_corner, hi_corner, node_type def _read_local_Header(header_file, dim): with open(header_file, "r") as f: t_snapshot = float(f.readline()) if dim==2: nx, nz = [int(x) for x in f.readline().split()] ny = 1 xmin, zmin = [float(x) for x in f.readline().split()] ymin = 0 xmax, zmax = [float(x) for x in f.readline().split()] ymax = 0 if dim==3: nx, ny, nz = [int(x) for x in f.readline().split()] xmin, ymin, zmin = [float(x) for x in f.readline().split()] xmax, ymax, zmax = [float(x) for x in f.readline().split()] field_names = f.readline().split() local_info = { 't_snapshot' : t_snapshot, 'field_names' : field_names, 'xmin' : xmin, 'ymin' : ymin, 'zmin' : zmin, 'xmax' : xmax, 'ymax' : ymax, 'zmax' : zmax, 'nx' : nx, 'ny' : ny, 'nz' : nz } return local_info ## ------------------------------------------------------------ ## USE THIS INSTEAD OF THE PREVIOUS FUNCTION IF Header contains ## (x,y,z) min and max vectors instead of zmin and zmax ## ------------------------------------------------------------ # def _read_local_Header(header_file): # with open(header_file, "r") as f: # t_snapshot = float(f.readline()) # axes_lo = [float(x) for x in f.readline().split()] # axes_hi = [float(x) for x in f.readline().split()] # local_info = { # 't_snapshot' : t_snapshot, # 'axes_lo' : axes_lo, # 'axes_hi' : axes_hi # } # return local_info def _read_global_Header(header_file): with open(header_file, "r") as f: nshapshots = int(f.readline()) dt_between_snapshots = float(f.readline()) gamma_boost = float(f.readline()) beta_boost = float(f.readline()) global_info = { 'nshapshots' : nshapshots, 'dt_between_snapshots' : dt_between_snapshots, 'gamma_boost' : gamma_boost, 'beta_boost' : beta_boost } return global_info def _read_header(header_file): with open(header_file, "r") as f: version = int(f.readline()) how = int(f.readline()) ncomp = int(f.readline()) # If the number of ghost cells is the same in all directions, # s is a string of the form '16\n'. # If the number of ghost cells varies depending on the direction, # s is a string of the form '(9,8)\n' in 2D or '(9,8,9)\n' in 3D. s = f.readline() s = s.replace('(', '') # remove left parenthesis '(', if any s = s.replace(')', '') # remove right parenthesis ')', if any nghost = np.fromstring(s, dtype = int, sep = ',') # convert from string to numpy array header = HeaderInfo(version, how, ncomp, nghost) # skip the next line f.readline() # read boxes boxes = [] for line in f: clean_line = line.strip().split() if clean_line == [')']: break lo_corner, hi_corner, node_type = _line_to_numpy_arrays(clean_line) boxes.append((lo_corner - nghost, hi_corner + nghost, node_type)) # read the file and offset position for the corresponding box file_names = [] offsets = [] for line in f: if line.startswith("FabOnDisk:"): clean_line = line.strip().split() file_names.append(clean_line[1]) offsets.append(int(clean_line[2])) return boxes, file_names, offsets, header def _combine_boxes(boxes): lo_corners, hi_corners = zip(*[(box[0], box[1]) for box in boxes]) domain_lo = np.min(lo_corners, axis=0) domain_hi = np.max(hi_corners, axis=0) return domain_lo, domain_hi def _read_field(raw_file, field_name): header_file = raw_file + field_name + "_H" boxes, file_names, offsets, header = _read_header(header_file) ng = header.nghost dom_lo, dom_hi = _combine_boxes(boxes) data_shape = dom_hi - dom_lo + 1 if header.ncomp > 1: data_shape = np.append(data_shape, header.ncomp) data = np.zeros(data_shape) for box, fn, offset in zip(boxes, file_names, offsets): lo = box[0] - dom_lo hi = box[1] - dom_lo shape = hi - lo + 1 if header.ncomp > 1: shape = np.append(shape, header.ncomp) with open(raw_file + fn, "rb") as f: f.seek(offset) if (header.version == 1): f.readline() # skip the first line arr = np.fromfile(f, 'float64', np.product(shape)) arr = arr.reshape(shape, order='F') box_shape = [slice(l,h+1) for l, h in zip(lo, hi)] if header.ncomp > 1: box_shape += [slice(None)] data[tuple(box_shape)] = arr return data def _read_buffer(snapshot, header_fn, _component_names): 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 shape = hi - lo + 1 size = np.product(shape) with open(snapshot + "/Level_0/" + fn, "rb") as f: f.seek(offset) if (header.version == 1): f.readline() # skip the first line arr = np.fromfile(f, 'float64', header.ncomp*size) for i in range(header.ncomp): comp_data = arr[i*size:(i+1)*size].reshape(shape, order='F') data = all_data[_component_names[i]] data[tuple([slice(l,h+1) for l, h in zip(lo, hi)])] = comp_data all_data[_component_names[i]] = data return all_data def read_reduced_diags(filename, delimiter=' '): ''' Read data written by WarpX Reduced Diagnostics, and return them into Python objects input: - filename name of file to open - delimiter (optional, default ',') delimiter between fields in header. output: - metadata_dict dictionary where first key is the type of metadata, second is the field - data dictionary with data ''' # Read header line unformatted_header = list( np.genfromtxt( filename, comments="@", max_rows=1, dtype="str", delimiter=delimiter) ) # From header line, get field name, units and column number field_names = [s[s.find("]")+1:s.find("(")] for s in unformatted_header] field_units = [s[s.find("(")+1:s.find(")")] for s in unformatted_header] field_column = [s[s.find("[")+1:s.find("]")] for s in unformatted_header] # Load data and re-format to a dictionary data = np.loadtxt( filename, delimiter=delimiter ) if data.ndim == 1: data_dict = {key: np.atleast_1d(data[i]) for i, key in enumerate(field_names)} else: data_dict = {key: data[:,i] for i, key in enumerate(field_names)} # Put header data into a dictionary metadata_dict = {} metadata_dict['units'] = {key: field_units[i] for i, key in enumerate(field_names)} metadata_dict['column'] = {key: field_column[i] for i, key in enumerate(field_names)} return metadata_dict, data_dict def read_reduced_diags_histogram(filename, delimiter=' '): ''' Modified based on read_reduced_diags Two extra return objects: - bin_value: the values of bins - bin_data: the histogram data values of bins ''' # Read header line unformatted_header = list( np.genfromtxt( filename, comments="@", max_rows=1, dtype="str", delimiter=delimiter) ) # From header line, get field name, units and column number field_names = [s[s.find("]")+1:s.find("(")] for s in unformatted_header] field_names[2:] = [s[s.find("b"):s.find("=")] for s in field_names[2:]] field_units = [s[s.find("(")+1:s.find(")")] for s in unformatted_header] field_column = [s[s.find("[")+1:s.find("]")] for s in unformatted_header] field_bin = [s[s.find("=")+1:s.find("(")] for s in unformatted_header] # Load data and re-format to a dictionary data = np.loadtxt( filename, delimiter=delimiter ) if data.ndim == 1: data_dict = {key: data[i] for i, key in enumerate(field_names)} else: data_dict = {key: data[:,i] for i, key in enumerate(field_names)} # Put header data into a dictionary metadata_dict = {} metadata_dict['units'] = {key: field_units[i] for i, key in enumerate(field_names)} metadata_dict['column'] = {key: field_column[i] for i, key in enumerate(field_names)} # Save bin values bin_value = np.asarray(field_bin[2:], dtype=np.float64, order='C') if data.ndim == 1: bin_data = data[2:] else: bin_data = data[:,2:] return metadata_dict, data_dict, bin_value, bin_data if __name__ == "__main__": data = read_lab_snapshot("lab_frame_data/snapshot00012", "lab_frame_data/Header");