diff --git a/astro_tigress/athena_read.py b/astro_tigress/athena_read.py index 8df309e..a183e56 100644 --- a/astro_tigress/athena_read.py +++ b/astro_tigress/athena_read.py @@ -17,6 +17,7 @@ # ======================================================================================== + def check_nan(data): """Check input NumPy array for the presence of any NaN entries""" if np.isnan(data).any(): @@ -26,12 +27,15 @@ def check_nan(data): # ======================================================================================== + def error_dat(filename, **kwargs): """Wrapper to np.loadtxt() for applying optional checks used in regression tests""" - data = np.loadtxt(filename, - dtype=np.float64, - ndmin=2, # prevent NumPy from squeezing singleton dimensions - **kwargs) + data = np.loadtxt( + filename, + dtype=np.float64, + ndmin=2, # prevent NumPy from squeezing singleton dimensions + **kwargs, + ) if check_nan_flag: check_nan(data) return data @@ -39,6 +43,7 @@ def error_dat(filename, **kwargs): # ======================================================================================== + def hst(filename, raw=False): """Read .hst files and return dict of 1D arrays. @@ -48,14 +53,14 @@ def hst(filename, raw=False): """ # Read data - with open(filename, 'r') as data_file: + with open(filename, "r") as data_file: # Find header header_found = False multiple_headers = False header_location = None line = data_file.readline() while len(line) > 0: - if line.startswith('# Athena'): + if line.startswith("# Athena"): if header_found: multiple_headers = True else: @@ -63,16 +68,18 @@ def hst(filename, raw=False): header_location = data_file.tell() line = data_file.readline() if multiple_headers: - warnings.warn('Multiple headers found; using most recent data', AthenaWarning) + warnings.warn( + "Multiple headers found; using most recent data", AthenaWarning + ) if header_location is None: - raise AthenaError('No header found') + raise AthenaError("No header found") # Parse header data_file.seek(header_location) header = data_file.readline() - data_names = re.findall(r'\[\d+\]=(\S+)', header) + data_names = re.findall(r"\[\d+\]=(\S+)", header) if len(data_names) == 0: - raise AthenaError('Header could not be parsed') + raise AthenaError("Header could not be parsed") # skip comments start_data = data_file.tell() line = data_file.readline() @@ -95,15 +102,17 @@ def hst(filename, raw=False): for key, val in list(data.items()): data[key] = np.array(val) if not raw: - if data_names[0] != 'time': - raise AthenaError('Cannot remove spurious data because time column could not' - + ' be identified') + if data_names[0] != "time": + raise AthenaError( + "Cannot remove spurious data because time column could not" + + " be identified" + ) branches_removed = False while not branches_removed: branches_removed = True - for n in range(1, len(data['time'])): - if data['time'][n] <= data['time'][n-1]: - branch_index = np.where(data['time'][:n] >= data['time'][n])[0][0] + for n in range(1, len(data["time"])): + if data["time"][n] <= data["time"][n - 1]: + branch_index = np.where(data["time"][:n] >= data["time"][n])[0][0] for key, val in list(data.items()): data[key] = np.concatenate((val[:branch_index], val[n:])) branches_removed = False @@ -116,6 +125,7 @@ def hst(filename, raw=False): # ======================================================================================== + def tab(filename, raw=False, dimensions=None): """Read .tab files and return dict or array. @@ -126,42 +136,44 @@ def tab(filename, raw=False, dimensions=None): # Check for valid number of dimensions if raw and not (dimensions == 1 or dimensions == 2 or dimensions == 3): - raise AthenaError('Improper number of dimensions') + raise AthenaError("Improper number of dimensions") if not raw and dimensions is not None: - warnings.warn('Ignoring specified number of dimensions', AthenaWarning) + warnings.warn("Ignoring specified number of dimensions", AthenaWarning) # Parse header if not raw: data_dict = {} - with open(filename, 'r') as data_file: + with open(filename, "r") as data_file: line = data_file.readline() - attributes = re.search(r'time=(\S+)\s+cycle=(\S+)\s+variables=(\S+)', line) + attributes = re.search(r"time=(\S+)\s+cycle=(\S+)\s+variables=(\S+)", line) line = data_file.readline() headings = line.split()[1:] - data_dict['time'] = float(attributes.group(1)) - data_dict['cycle'] = int(attributes.group(2)) - data_dict['variables'] = attributes.group(3) - if headings[0] == 'i' and headings[2] == 'j' and headings[4] == 'k': + data_dict["time"] = float(attributes.group(1)) + data_dict["cycle"] = int(attributes.group(2)) + data_dict["variables"] = attributes.group(3) + if headings[0] == "i" and headings[2] == "j" and headings[4] == "k": headings = headings[1:2] + headings[3:4] + headings[5:] dimensions = 3 - elif ((headings[0] == 'i' and headings[2] == 'j') - or (headings[0] == 'i' and headings[2] == 'k') - or (headings[0] == 'j' and headings[2] == 'k')): + elif ( + (headings[0] == "i" and headings[2] == "j") + or (headings[0] == "i" and headings[2] == "k") + or (headings[0] == "j" and headings[2] == "k") + ): headings = headings[1:2] + headings[3:] dimensions = 2 - elif headings[0] == 'i' or headings[0] == 'j' or headings[0] == 'k': + elif headings[0] == "i" or headings[0] == "j" or headings[0] == "k": headings = headings[1:] dimensions = 1 else: - raise AthenaError('Could not parse header') + raise AthenaError("Could not parse header") # Go through lines data_array = [] - with open(filename, 'r') as data_file: + with open(filename, "r") as data_file: first_line = True for line in data_file: # Skip comments - if line.split()[0][0] == '#': + if line.split()[0][0] == "#": continue # Extract cell indices @@ -192,13 +204,18 @@ def tab(filename, raw=False, dimensions=None): # Reshape array based on number of dimensions if dimensions == 1: - array_shape = (i_max-i_min+1, num_entries) + array_shape = (i_max - i_min + 1, num_entries) array_transpose = (1, 0) if dimensions == 2: - array_shape = (j_max-j_min+1, i_max-i_min+1, num_entries) + array_shape = (j_max - j_min + 1, i_max - i_min + 1, num_entries) array_transpose = (2, 0, 1) if dimensions == 3: - array_shape = (k_max-k_min+1, j_max-j_min+1, i_max-i_min+1, num_entries) + array_shape = ( + k_max - k_min + 1, + j_max - j_min + 1, + i_max - i_min + 1, + num_entries, + ) array_transpose = (3, 0, 1, 2) data_array = np.transpose(np.reshape(data_array, array_shape), array_transpose) @@ -217,19 +234,20 @@ def tab(filename, raw=False, dimensions=None): # ======================================================================================== + def vtk(filename): """Read .vtk files and return dict of arrays of data.""" # Read raw data - with open(filename, 'rb') as data_file: + with open(filename, "rb") as data_file: raw_data = data_file.read() - raw_data_ascii = raw_data.decode('ascii', 'replace') + raw_data_ascii = raw_data.decode("ascii", "replace") # Skip header current_index = 0 current_char = raw_data_ascii[current_index] - while current_char == '#': - while current_char != '\n': + while current_char == "#": + while current_char != "\n": current_index += 1 current_char = raw_data_ascii[current_index] current_index += 1 @@ -238,90 +256,96 @@ def vtk(filename): # Function for skipping though the file def skip_string(expected_string): expected_string_len = len(expected_string) - if (raw_data_ascii[current_index:current_index+expected_string_len] - != expected_string): - raise AthenaError('File not formatted as expected') - return current_index+expected_string_len + if ( + raw_data_ascii[current_index : current_index + expected_string_len] + != expected_string + ): + raise AthenaError("File not formatted as expected") + return current_index + expected_string_len # Read metadata - current_index = skip_string('BINARY\nDATASET RECTILINEAR_GRID\nDIMENSIONS ') + current_index = skip_string("BINARY\nDATASET RECTILINEAR_GRID\nDIMENSIONS ") end_of_line_index = current_index + 1 - while raw_data_ascii[end_of_line_index] != '\n': + while raw_data_ascii[end_of_line_index] != "\n": end_of_line_index += 1 data_to_map = raw_data_ascii[current_index:end_of_line_index] - face_dimensions = list(map(int, data_to_map.split(' '))) + face_dimensions = list(map(int, data_to_map.split(" "))) current_index = end_of_line_index + 1 # Function for reading interface locations def read_faces(letter, num_faces): - identifier_string = '{0}_COORDINATES {1} float\n'.format(letter, num_faces) + identifier_string = "{0}_COORDINATES {1} float\n".format(letter, num_faces) begin_index = skip_string(identifier_string) - format_string = '>' + 'f'*num_faces - end_index = begin_index + 4*num_faces + format_string = ">" + "f" * num_faces + end_index = begin_index + 4 * num_faces vals = np.array(struct.unpack(format_string, raw_data[begin_index:end_index])) - return vals, end_index+1 + return vals, end_index + 1 # Read interface locations - x_faces, current_index = read_faces('X', face_dimensions[0]) - y_faces, current_index = read_faces('Y', face_dimensions[1]) - z_faces, current_index = read_faces('Z', face_dimensions[2]) + x_faces, current_index = read_faces("X", face_dimensions[0]) + y_faces, current_index = read_faces("Y", face_dimensions[1]) + z_faces, current_index = read_faces("Z", face_dimensions[2]) # Prepare to read quantities defined on grid - cell_dimensions = np.array([max(dim-1, 1) for dim in face_dimensions]) + cell_dimensions = np.array([max(dim - 1, 1) for dim in face_dimensions]) num_cells = cell_dimensions.prod() - current_index = skip_string('CELL_DATA {0}\n'.format(num_cells)) - if raw_data_ascii[current_index:current_index+1] == '\n': - current_index = skip_string('\n') # extra newline inserted by join script + current_index = skip_string("CELL_DATA {0}\n".format(num_cells)) + if raw_data_ascii[current_index : current_index + 1] == "\n": + current_index = skip_string("\n") # extra newline inserted by join script data = {} # Function for reading scalar data def read_cell_scalars(): - begin_index = skip_string('SCALARS ') + begin_index = skip_string("SCALARS ") end_of_word_index = begin_index + 1 - while raw_data_ascii[end_of_word_index] != ' ': + while raw_data_ascii[end_of_word_index] != " ": end_of_word_index += 1 array_name = raw_data_ascii[begin_index:end_of_word_index] - string_to_skip = 'SCALARS {0} float\nLOOKUP_TABLE default\n'.format(array_name) + string_to_skip = "SCALARS {0} float\nLOOKUP_TABLE default\n".format(array_name) begin_index = skip_string(string_to_skip) - format_string = '>' + 'f'*num_cells - end_index = begin_index + 4*num_cells + format_string = ">" + "f" * num_cells + end_index = begin_index + 4 * num_cells data[array_name] = struct.unpack(format_string, raw_data[begin_index:end_index]) dimensions = tuple(cell_dimensions[::-1]) data[array_name] = np.array(data[array_name]).reshape(dimensions) - return end_index+1 + return end_index + 1 # Function for reading vector data def read_cell_vectors(): - begin_index = skip_string('VECTORS ') + begin_index = skip_string("VECTORS ") end_of_word_index = begin_index + 1 - while raw_data_ascii[end_of_word_index] != '\n': + while raw_data_ascii[end_of_word_index] != "\n": end_of_word_index += 1 array_name = raw_data_ascii[begin_index:end_of_word_index] - string_to_skip = 'VECTORS {0}\n'.format(array_name) + string_to_skip = "VECTORS {0}\n".format(array_name) array_name = array_name[:-6] # remove ' float' begin_index = skip_string(string_to_skip) - format_string = '>' + 'f'*num_cells*3 - end_index = begin_index + 4*num_cells*3 + format_string = ">" + "f" * num_cells * 3 + end_index = begin_index + 4 * num_cells * 3 data[array_name] = struct.unpack(format_string, raw_data[begin_index:end_index]) dimensions = tuple(np.append(cell_dimensions[::-1], 3)) data[array_name] = np.array(data[array_name]).reshape(dimensions) - return end_index+1 + return end_index + 1 # Read quantities defined on grid while current_index < len(raw_data): - expected_string = 'SCALARS' + expected_string = "SCALARS" expected_string_len = len(expected_string) - if (raw_data_ascii[current_index:current_index+expected_string_len] - == expected_string): + if ( + raw_data_ascii[current_index : current_index + expected_string_len] + == expected_string + ): current_index = read_cell_scalars() continue - expected_string = 'VECTORS' + expected_string = "VECTORS" expected_string_len = len(expected_string) - if (raw_data_ascii[current_index:current_index+expected_string_len] - == expected_string): + if ( + raw_data_ascii[current_index : current_index + expected_string_len] + == expected_string + ): current_index = read_cell_vectors() continue - raise AthenaError('File not formatted as expected') + raise AthenaError("File not formatted as expected") if check_nan_flag: check_nan(x_faces) @@ -335,11 +359,33 @@ def read_cell_vectors(): # ======================================================================================== -def athdf(filename, raw=False, data=None, quantities=None, dtype=None, level=None, - return_levels=False, subsample=False, fast_restrict=False, x1_min=None, - x1_max=None, x2_min=None, x2_max=None, x3_min=None, x3_max=None, vol_func=None, - vol_params=None, face_func_1=None, face_func_2=None, face_func_3=None, - center_func_1=None, center_func_2=None, center_func_3=None, num_ghost=0): + +def athdf( + filename, + raw=False, + data=None, + quantities=None, + dtype=None, + level=None, + return_levels=False, + subsample=False, + fast_restrict=False, + x1_min=None, + x1_max=None, + x2_min=None, + x2_max=None, + x3_min=None, + x3_max=None, + vol_func=None, + vol_params=None, + face_func_1=None, + face_func_2=None, + face_func_3=None, + center_func_1=None, + center_func_2=None, + center_func_3=None, + num_ghost=0, +): """Read .athdf files and populate dict of arrays of data. @@ -353,30 +399,32 @@ def athdf(filename, raw=False, data=None, quantities=None, dtype=None, level=Non # Handle request for raw data if raw: # Open file - with h5py.File(filename, 'r') as f: + with h5py.File(filename, "r") as f: # Store file-level attributes data = {} for key in f.attrs: data[str(key)] = f.attrs[key] # Store location metadata - data['Levels'] = f['Levels'][:] - data['LogicalLocations'] = f['LogicalLocations'][:] + data["Levels"] = f["Levels"][:] + data["LogicalLocations"] = f["LogicalLocations"][:] # Store coordinate data - data['x1f'] = f['x1f'][:] - data['x2f'] = f['x2f'][:] - data['x3f'] = f['x3f'][:] - data['x1v'] = f['x1v'][:] - data['x2v'] = f['x2v'][:] - data['x3v'] = f['x3v'][:] + data["x1f"] = f["x1f"][:] + data["x2f"] = f["x2f"][:] + data["x3f"] = f["x3f"][:] + data["x1v"] = f["x1v"][:] + data["x2v"] = f["x2v"][:] + data["x3v"] = f["x3v"][:] # Get metadata describing file layout - dataset_names = np.array([x.decode('ascii', 'replace') - for x in f.attrs['DatasetNames'][:]]) - dataset_sizes = f.attrs['NumVariables'][:] - variable_names = np.array([x.decode('ascii', 'replace') - for x in f.attrs['VariableNames'][:]]) + dataset_names = np.array( + [x.decode("ascii", "replace") for x in f.attrs["DatasetNames"][:]] + ) + dataset_sizes = f.attrs["NumVariables"][:] + variable_names = np.array( + [x.decode("ascii", "replace") for x in f.attrs["VariableNames"][:]] + ) # Store cell data for dataset_index, dataset_name in enumerate(dataset_names): @@ -401,39 +449,50 @@ def athdf(filename, raw=False, data=None, quantities=None, dtype=None, level=Non new_data = False # Open file - with h5py.File(filename, 'r') as f: + with h5py.File(filename, "r") as f: # Extract size information - max_level = f.attrs['MaxLevel'] + max_level = f.attrs["MaxLevel"] if level is None: level = max_level - block_size = f.attrs['MeshBlockSize'] - root_grid_size = f.attrs['RootGridSize'] - levels = f['Levels'][:] - logical_locations = f['LogicalLocations'][:] + block_size = f.attrs["MeshBlockSize"] + root_grid_size = f.attrs["RootGridSize"] + levels = f["Levels"][:] + logical_locations = f["LogicalLocations"][:] if dtype is None: - dtype = f[f.attrs['DatasetNames'][0]].dtype.newbyteorder('=') - if num_ghost == 0 and np.array(f['x1v']).min() < f.attrs['RootGridX1'][0]: - raise AthenaError('Ghost zones detected but "num_ghost" keyword set to zero.') + dtype = f[f.attrs["DatasetNames"][0]].dtype.newbyteorder("=") + if num_ghost == 0 and np.array(f["x1v"]).min() < f.attrs["RootGridX1"][0]: + raise AthenaError( + 'Ghost zones detected but "num_ghost" keyword set to zero.' + ) if num_ghost > 0 and not np.all(levels == max_level): - raise AthenaError('Cannot use ghost zones with different refinement levels') + raise AthenaError("Cannot use ghost zones with different refinement levels") nx_vals = [] for d in range(3): if block_size[d] == 1 and root_grid_size[d] > 1: # sum or slice - other_locations = [location - for location in zip(levels, - logical_locations[:, (d+1) % 3], - logical_locations[:, (d+2) % 3])] + other_locations = [ + location + for location in zip( + levels, + logical_locations[:, (d + 1) % 3], + logical_locations[:, (d + 2) % 3], + ) + ] if len(set(other_locations)) == len(other_locations): # effective slice nx_vals.append(1) else: # nontrivial sum num_blocks_this_dim = 0 - for level_this_dim, loc_this_dim in zip(levels, - logical_locations[:, d]): + for level_this_dim, loc_this_dim in zip( + levels, logical_locations[:, d] + ): if level_this_dim <= level: - possible_max = (loc_this_dim+1) * 2**(level-level_this_dim) + possible_max = (loc_this_dim + 1) * 2 ** ( + level - level_this_dim + ) num_blocks_this_dim = max(num_blocks_this_dim, possible_max) else: - possible_max = (loc_this_dim+1) / 2**(level_this_dim-level) + possible_max = (loc_this_dim + 1) / 2 ** ( + level_this_dim - level + ) num_blocks_this_dim = max(num_blocks_this_dim, possible_max) nx_vals.append(num_blocks_this_dim) elif block_size[d] == 1: # singleton dimension @@ -452,34 +511,55 @@ def athdf(filename, raw=False, data=None, quantities=None, dtype=None, level=Non num_extended_dims += 1 # Set volume function for preset coordinates if needed - coord = f.attrs['Coordinates'].decode('ascii', 'replace') - if level < max_level and not subsample and not fast_restrict and vol_func is None: - x1_rat = f.attrs['RootGridX1'][2] - x2_rat = f.attrs['RootGridX2'][2] - x3_rat = f.attrs['RootGridX3'][2] - if (coord == 'cartesian' or coord == 'minkowski' or coord == 'tilted' - or coord == 'sinusoidal'): - if ((nx1 == 1 or x1_rat == 1.0) and (nx2 == 1 or x2_rat == 1.0) - and (nx3 == 1 or x3_rat == 1.0)): + coord = f.attrs["Coordinates"].decode("ascii", "replace") + if ( + level < max_level + and not subsample + and not fast_restrict + and vol_func is None + ): + x1_rat = f.attrs["RootGridX1"][2] + x2_rat = f.attrs["RootGridX2"][2] + x3_rat = f.attrs["RootGridX3"][2] + if ( + coord == "cartesian" + or coord == "minkowski" + or coord == "tilted" + or coord == "sinusoidal" + ): + if ( + (nx1 == 1 or x1_rat == 1.0) + and (nx2 == 1 or x2_rat == 1.0) + and (nx3 == 1 or x3_rat == 1.0) + ): fast_restrict = True else: + def vol_func(xm, xp, ym, yp, zm, zp): - return (xp-xm) * (yp-ym) * (zp-zm) - elif coord == 'cylindrical': - if (nx1 == 1 and (nx2 == 1 or x2_rat == 1.0) - and (nx3 == 1 or x3_rat == 1.0)): + return (xp - xm) * (yp - ym) * (zp - zm) + elif coord == "cylindrical": + if ( + nx1 == 1 + and (nx2 == 1 or x2_rat == 1.0) + and (nx3 == 1 or x3_rat == 1.0) + ): fast_restrict = True else: + def vol_func(rm, rp, phim, phip, zm, zp): - return (rp**2-rm**2) * (phip-phim) * (zp-zm) - elif coord == 'spherical_polar' or coord == 'schwarzschild': + return (rp**2 - rm**2) * (phip - phim) * (zp - zm) + elif coord == "spherical_polar" or coord == "schwarzschild": if nx1 == 1 and nx2 == 1 and (nx3 == 1 or x3_rat == 1.0): fast_restrict = True else: + def vol_func(rm, rp, thetam, thetap, phim, phip): - return ((rp**3-rm**3) * abs(np.cos(thetam)-np.cos(thetap)) - * (phip-phim)) - elif coord == 'kerr-schild': + return ( + (rp**3 - rm**3) + * abs(np.cos(thetam) - np.cos(thetap)) + * (phip - phim) + ) + elif coord == "kerr-schild": if nx1 == 1 and nx2 == 1 and (nx3 == 1 or x3_rat == 1.0): fast_restrict = True else: @@ -488,86 +568,124 @@ def vol_func(rm, rp, thetam, thetap, phim, phip): def vol_func(rm, rp, thetam, thetap, phim, phip): cosm = np.cos(thetam) cosp = np.cos(thetap) - return (((rp**3-rm**3) * abs(cosm-cosp) - + a**2 * (rp-rm) * abs(cosm**3-cosp**3)) * (phip-phim)) + return ( + (rp**3 - rm**3) * abs(cosm - cosp) + + a**2 * (rp - rm) * abs(cosm**3 - cosp**3) + ) * (phip - phim) else: - raise AthenaError('Coordinates not recognized') + raise AthenaError("Coordinates not recognized") # Set cell center functions for preset coordinates if center_func_1 is None: - if (coord == 'cartesian' or coord == 'minkowski' or coord == 'tilted' - or coord == 'sinusoidal' or coord == 'kerr-schild'): + if ( + coord == "cartesian" + or coord == "minkowski" + or coord == "tilted" + or coord == "sinusoidal" + or coord == "kerr-schild" + ): + def center_func_1(xm, xp): - return 0.5 * (xm+xp) - elif coord == 'cylindrical': + return 0.5 * (xm + xp) + elif coord == "cylindrical": + def center_func_1(xm, xp): - return 2.0/3.0 * (xp**3-xm**3) / (xp**2-xm**2) - elif coord == 'spherical_polar': + return 2.0 / 3.0 * (xp**3 - xm**3) / (xp**2 - xm**2) + elif coord == "spherical_polar": + def center_func_1(xm, xp): - return 3.0/4.0 * (xp**4-xm**4) / (xp**3-xm**3) - elif coord == 'schwarzschild': + return 3.0 / 4.0 * (xp**4 - xm**4) / (xp**3 - xm**3) + elif coord == "schwarzschild": + def center_func_1(xm, xp): - return (0.5*(xm**3+xp**3)) ** 1.0/3.0 + return (0.5 * (xm**3 + xp**3)) ** 1.0 / 3.0 else: - raise AthenaError('Coordinates not recognized') + raise AthenaError("Coordinates not recognized") if center_func_2 is None: - if (coord == 'cartesian' or coord == 'cylindrical' or coord == 'minkowski' - or coord == 'tilted' or coord == 'sinusoidal' - or coord == 'kerr-schild'): + if ( + coord == "cartesian" + or coord == "cylindrical" + or coord == "minkowski" + or coord == "tilted" + or coord == "sinusoidal" + or coord == "kerr-schild" + ): + def center_func_2(xm, xp): - return 0.5 * (xm+xp) - elif coord == 'spherical_polar': + return 0.5 * (xm + xp) + elif coord == "spherical_polar": + def center_func_2(xm, xp): sm = np.sin(xm) cm = np.cos(xm) sp = np.sin(xp) cp = np.cos(xp) - return (sp-xp*cp - sm+xm*cm) / (cm - cp) - elif coord == 'schwarzschild': + return (sp - xp * cp - sm + xm * cm) / (cm - cp) + elif coord == "schwarzschild": + def center_func_2(xm, xp): return np.arccos(0.5 * (np.cos(xm) + np.cos(xp))) else: - raise AthenaError('Coordinates not recognized') + raise AthenaError("Coordinates not recognized") if center_func_3 is None: - if (coord == 'cartesian' or coord == 'cylindrical' or coord == 'tilted' - or coord == 'spherical_polar' or coord == 'minkowski' - or coord == 'sinusoidal' or coord == 'schwarzschild' - or coord == 'kerr-schild'): + if ( + coord == "cartesian" + or coord == "cylindrical" + or coord == "tilted" + or coord == "spherical_polar" + or coord == "minkowski" + or coord == "sinusoidal" + or coord == "schwarzschild" + or coord == "kerr-schild" + ): def center_func_3(xm, xp): - return 0.5 * (xm+xp) + return 0.5 * (xm + xp) else: - raise AthenaError('Coordinates not recognized') + raise AthenaError("Coordinates not recognized") # Check output level compared to max level in file if level < max_level and not subsample and not fast_restrict: - warnings.warn('Exact restriction being used: performance severely affected;' - + ' see documentation', AthenaWarning) + warnings.warn( + "Exact restriction being used: performance severely affected;" + + " see documentation", + AthenaWarning, + ) sys.stderr.flush() if level > max_level: - warnings.warn('Requested refinement level higher than maximum level in file:' - + ' all cells will be prolongated', AthenaWarning) + warnings.warn( + "Requested refinement level higher than maximum level in file:" + + " all cells will be prolongated", + AthenaWarning, + ) sys.stderr.flush() # Check that subsampling and/or fast restriction will work if needed if level < max_level and (subsample or fast_restrict): if num_ghost > 0: - raise AthenaError('Subsampling and fast restriction incompatible with' - + ' ghost zones') - max_restrict_factor = 2**(max_level-level) + raise AthenaError( + "Subsampling and fast restriction incompatible with" + + " ghost zones" + ) + max_restrict_factor = 2 ** (max_level - level) for current_block_size in block_size: - if (current_block_size != 1 - and current_block_size % max_restrict_factor != 0): - raise AthenaError('Block boundaries at finest level must be cell' - + ' boundaries at desired level for subsampling or' - + ' fast restriction to work') + if ( + current_block_size != 1 + and current_block_size % max_restrict_factor != 0 + ): + raise AthenaError( + "Block boundaries at finest level must be cell" + + " boundaries at desired level for subsampling or" + + " fast restriction to work" + ) # Create list of all quantities if none given - var_quantities = np.array([x.decode('ascii', 'replace') - for x in f.attrs['VariableNames'][:]]) - coord_quantities = ('x1f', 'x2f', 'x3f', 'x1v', 'x2v', 'x3v') + var_quantities = np.array( + [x.decode("ascii", "replace") for x in f.attrs["VariableNames"][:]] + ) + coord_quantities = ("x1f", "x2f", "x3f", "x1v", "x2v", "x3v") attr_quantities = [key for key in f.attrs] - other_quantities = ('Levels',) + other_quantities = ("Levels",) if not new_data: quantities = list(data.values()) elif quantities is None: @@ -577,24 +695,33 @@ def center_func_3(xm, xp): if q not in var_quantities and q not in coord_quantities: possibilities = '", "'.join(var_quantities) possibilities = '"' + possibilities + '"' - error_string = ('Quantity not recognized: file does not include "{0}"' - + ' but does include {1}') + error_string = ( + 'Quantity not recognized: file does not include "{0}"' + + " but does include {1}" + ) raise AthenaError(error_string.format(q, possibilities)) - quantities = [str(q) for q in quantities if q not in coord_quantities - and q not in attr_quantities and q not in other_quantities] + quantities = [ + str(q) + for q in quantities + if q not in coord_quantities + and q not in attr_quantities + and q not in other_quantities + ] # Store file attribute metadata for key in attr_quantities: data[str(key)] = f.attrs[key] # Get metadata describing file layout - num_blocks = f.attrs['NumMeshBlocks'] - dataset_names = np.array([x.decode('ascii', 'replace') - for x in f.attrs['DatasetNames'][:]]) - dataset_sizes = f.attrs['NumVariables'][:] + num_blocks = f.attrs["NumMeshBlocks"] + dataset_names = np.array( + [x.decode("ascii", "replace") for x in f.attrs["DatasetNames"][:]] + ) + dataset_sizes = f.attrs["NumVariables"][:] dataset_sizes_cumulative = np.cumsum(dataset_sizes) - variable_names = np.array([x.decode('ascii', 'replace') - for x in f.attrs['VariableNames'][:]]) + variable_names = np.array( + [x.decode("ascii", "replace") for x in f.attrs["VariableNames"][:]] + ) quantity_datasets = [] quantity_indices = [] for q in quantities: @@ -603,66 +730,81 @@ def center_func_3(xm, xp): if dataset_num == 0: dataset_index = var_num else: - dataset_index = var_num - dataset_sizes_cumulative[dataset_num-1] + dataset_index = var_num - dataset_sizes_cumulative[dataset_num - 1] quantity_datasets.append(dataset_names[dataset_num]) quantity_indices.append(dataset_index) # Locate fine block for coordinates in case of slice fine_block = np.where(levels == max_level)[0][0] - x1m = f['x1f'][fine_block, 0] - x1p = f['x1f'][fine_block, 1] - x2m = f['x2f'][fine_block, 0] - x2p = f['x2f'][fine_block, 1] - x3m = f['x3f'][fine_block, 0] - x3p = f['x3f'][fine_block, 1] + x1m = f["x1f"][fine_block, 0] + x1p = f["x1f"][fine_block, 1] + x2m = f["x2f"][fine_block, 0] + x2p = f["x2f"][fine_block, 1] + x3m = f["x3f"][fine_block, 0] + x3p = f["x3f"][fine_block, 1] # Populate coordinate arrays face_funcs = (face_func_1, face_func_2, face_func_3) center_funcs = (center_func_1, center_func_2, center_func_3) - for d, nx, face_func, center_func in zip(list(range(1, 4)), nx_vals, face_funcs, - center_funcs): - xf = 'x' + repr(d) + 'f' - xv = 'x' + repr(d) + 'v' + for d, nx, face_func, center_func in zip( + list(range(1, 4)), nx_vals, face_funcs, center_funcs + ): + xf = "x" + repr(d) + "f" + xv = "x" + repr(d) + "v" if nx == 1: - xm = (x1m, x2m, x3m)[d-1] - xp = (x1p, x2p, x3p)[d-1] + xm = (x1m, x2m, x3m)[d - 1] + xp = (x1p, x2p, x3p)[d - 1] data[xf] = np.array([xm, xp], dtype=dtype) else: - xmin = f.attrs['RootGridX' + repr(d)][0] - xmax = f.attrs['RootGridX' + repr(d)][1] - xrat_root = f.attrs['RootGridX' + repr(d)][2] + xmin = f.attrs["RootGridX" + repr(d)][0] + xmax = f.attrs["RootGridX" + repr(d)][1] + xrat_root = f.attrs["RootGridX" + repr(d)][2] if xrat_root == -1.0 and face_func is None: - raise AthenaError('Must specify user-defined face_func_{0}'.format(d)) + raise AthenaError( + "Must specify user-defined face_func_{0}".format(d) + ) elif face_func is not None: if num_ghost > 0: - raise AthenaError('Ghost zones incompatible with user-defined' - + ' coordinate spacing') - data[xf] = face_func(xmin, xmax, xrat_root, nx+1) + raise AthenaError( + "Ghost zones incompatible with user-defined" + + " coordinate spacing" + ) + data[xf] = face_func(xmin, xmax, xrat_root, nx + 1) elif xrat_root == 1.0: if np.all(levels == level): data[xf] = np.empty(nx + 1, dtype=dtype) - for n_block in range(int((nx - 2*num_ghost) - / (block_size[d-1] - 2*num_ghost))): - sample_block = np.where(logical_locations[:, d-1] - == n_block)[0][0] - index_low = n_block * (block_size[d-1] - 2*num_ghost) - index_high = index_low + block_size[d-1] + 1 + for n_block in range( + int( + (nx - 2 * num_ghost) + / (block_size[d - 1] - 2 * num_ghost) + ) + ): + sample_block = np.where( + logical_locations[:, d - 1] == n_block + )[0][0] + index_low = n_block * (block_size[d - 1] - 2 * num_ghost) + index_high = index_low + block_size[d - 1] + 1 data[xf][index_low:index_high] = f[xf][sample_block, :] else: if num_ghost > 0: - raise AthenaError('Cannot use ghost zones with different' - + ' refinement levels') + raise AthenaError( + "Cannot use ghost zones with different" + + " refinement levels" + ) data[xf] = np.linspace(xmin, xmax, nx + 1, dtype=dtype) else: if num_ghost > 0: - raise AthenaError('Ghost zones incompatible with non-uniform' - + ' coordinate spacing') + raise AthenaError( + "Ghost zones incompatible with non-uniform" + + " coordinate spacing" + ) xrat = xrat_root ** (1.0 / 2**level) - data[xf] = (xmin + (1.0-xrat**np.arange(nx+1, dtype=dtype)) - / (1.0-xrat**nx) * (xmax-xmin)) + data[xf] = xmin + (1.0 - xrat ** np.arange(nx + 1, dtype=dtype)) / ( + 1.0 - xrat**nx + ) * (xmax - xmin) data[xv] = np.empty(nx, dtype=dtype) for i in range(nx): - data[xv][i] = center_func(data[xf][i], data[xf][i+1]) + data[xv][i] = center_func(data[xf][i], data[xf][i + 1]) # Account for selection x1_select = False @@ -672,61 +814,73 @@ def center_func_3(xm, xp): i_max = nx1 j_max = nx2 k_max = nx3 - error_string = '{0} must be {1} than {2} in order to intersect data range' - if x1_min is not None and x1_min >= data['x1f'][1]: - if x1_min >= data['x1f'][-1]: - raise AthenaError(error_string.format('x1_min', 'less', data['x1f'][-1])) + error_string = "{0} must be {1} than {2} in order to intersect data range" + if x1_min is not None and x1_min >= data["x1f"][1]: + if x1_min >= data["x1f"][-1]: + raise AthenaError( + error_string.format("x1_min", "less", data["x1f"][-1]) + ) x1_select = True - i_min = np.where(data['x1f'] <= x1_min)[0][-1] - if x1_max is not None and x1_max <= data['x1f'][-2]: - if x1_max <= data['x1f'][0]: - raise AthenaError(error_string.format('x1_max', 'greater', - data['x1f'][0])) + i_min = np.where(data["x1f"] <= x1_min)[0][-1] + if x1_max is not None and x1_max <= data["x1f"][-2]: + if x1_max <= data["x1f"][0]: + raise AthenaError( + error_string.format("x1_max", "greater", data["x1f"][0]) + ) x1_select = True - i_max = np.where(data['x1f'] >= x1_max)[0][0] - if x2_min is not None and x2_min >= data['x2f'][1]: - if x2_min >= data['x2f'][-1]: - raise AthenaError(error_string.format('x2_min', 'less', data['x2f'][-1])) + i_max = np.where(data["x1f"] >= x1_max)[0][0] + if x2_min is not None and x2_min >= data["x2f"][1]: + if x2_min >= data["x2f"][-1]: + raise AthenaError( + error_string.format("x2_min", "less", data["x2f"][-1]) + ) x2_select = True - j_min = np.where(data['x2f'] <= x2_min)[0][-1] - if x2_max is not None and x2_max <= data['x2f'][-2]: - if x2_max <= data['x2f'][0]: - raise AthenaError(error_string.format('x2_max', 'greater', - data['x2f'][0])) + j_min = np.where(data["x2f"] <= x2_min)[0][-1] + if x2_max is not None and x2_max <= data["x2f"][-2]: + if x2_max <= data["x2f"][0]: + raise AthenaError( + error_string.format("x2_max", "greater", data["x2f"][0]) + ) x2_select = True - j_max = np.where(data['x2f'] >= x2_max)[0][0] - if x3_min is not None and x3_min >= data['x3f'][1]: - if x3_min >= data['x3f'][-1]: - raise AthenaError(error_string.format('x3_min', 'less', data['x3f'][-1])) + j_max = np.where(data["x2f"] >= x2_max)[0][0] + if x3_min is not None and x3_min >= data["x3f"][1]: + if x3_min >= data["x3f"][-1]: + raise AthenaError( + error_string.format("x3_min", "less", data["x3f"][-1]) + ) x3_select = True - k_min = np.where(data['x3f'] <= x3_min)[0][-1] - if x3_max is not None and x3_max <= data['x3f'][-2]: - if x3_max <= data['x3f'][0]: - raise AthenaError(error_string.format('x3_max', 'greater', - data['x3f'][0])) + k_min = np.where(data["x3f"] <= x3_min)[0][-1] + if x3_max is not None and x3_max <= data["x3f"][-2]: + if x3_max <= data["x3f"][0]: + raise AthenaError( + error_string.format("x3_max", "greater", data["x3f"][0]) + ) x3_select = True - k_max = np.where(data['x3f'] >= x3_max)[0][0] + k_max = np.where(data["x3f"] >= x3_max)[0][0] if (x1_select or x2_select or x3_select) and num_ghost > 0: - raise AthenaError('Cannot take subsets with ghost zones') + raise AthenaError("Cannot take subsets with ghost zones") # Adjust coordinates if selection made if x1_select: - data['x1f'] = data['x1f'][i_min:i_max+1] - data['x1v'] = data['x1v'][i_min:i_max] + data["x1f"] = data["x1f"][i_min : i_max + 1] + data["x1v"] = data["x1v"][i_min:i_max] if x2_select: - data['x2f'] = data['x2f'][j_min:j_max+1] - data['x2v'] = data['x2v'][j_min:j_max] + data["x2f"] = data["x2f"][j_min : j_max + 1] + data["x2v"] = data["x2v"][j_min:j_max] if x3_select: - data['x3f'] = data['x3f'][k_min:k_max+1] - data['x3v'] = data['x3v'][k_min:k_max] + data["x3f"] = data["x3f"][k_min : k_max + 1] + data["x3v"] = data["x3v"][k_min:k_max] # Prepare arrays for data and bookkeeping if new_data: for q in quantities: - data[q] = np.zeros((k_max-k_min, j_max-j_min, i_max-i_min), dtype=dtype) + data[q] = np.zeros( + (k_max - k_min, j_max - j_min, i_max - i_min), dtype=dtype + ) if return_levels: - data['Levels'] = np.empty((k_max-k_min, j_max-j_min, i_max-i_min), - dtype=np.int32) + data["Levels"] = np.empty( + (k_max - k_min, j_max - j_min, i_max - i_min), dtype=np.int32 + ) else: for q in quantities: data[q].fill(0.0) @@ -745,12 +899,21 @@ def center_func_3(xm, xp): s = 2 ** (level - block_level) # Calculate destination indices, without selection - il_d = (block_location[0] * (block_size[0] - 2*num_ghost) * s - if nx1 > 1 else 0) - jl_d = (block_location[1] * (block_size[1] - 2*num_ghost) * s - if nx2 > 1 else 0) - kl_d = (block_location[2] * (block_size[2] - 2*num_ghost) * s - if nx3 > 1 else 0) + il_d = ( + block_location[0] * (block_size[0] - 2 * num_ghost) * s + if nx1 > 1 + else 0 + ) + jl_d = ( + block_location[1] * (block_size[1] - 2 * num_ghost) * s + if nx2 > 1 + else 0 + ) + kl_d = ( + block_location[2] * (block_size[2] - 2 * num_ghost) * s + if nx3 > 1 + else 0 + ) iu_d = il_d + block_size[0] * s if nx1 > 1 else 1 ju_d = jl_d + block_size[1] * s if nx2 > 1 else 1 ku_d = kl_d + block_size[2] * s if nx3 > 1 else 1 @@ -774,8 +937,9 @@ def center_func_3(xm, xp): ku_d = min(ku_d, k_max) - k_min # Assign values - for q, dataset, index in zip(quantities, quantity_datasets, - quantity_indices): + for q, dataset, index in zip( + quantities, quantity_datasets, quantity_indices + ): block_data = f[dataset][index, block_num, :] if s > 1: if nx1 > 1: @@ -784,9 +948,9 @@ def center_func_3(xm, xp): block_data = np.repeat(block_data, s, axis=1) if nx3 > 1: block_data = np.repeat(block_data, s, axis=0) - data[q][kl_d:ku_d, jl_d:ju_d, il_d:iu_d] = block_data[kl_s:ku_s, - jl_s:ju_s, - il_s:iu_s] + data[q][kl_d:ku_d, jl_d:ju_d, il_d:iu_d] = block_data[ + kl_s:ku_s, jl_s:ju_s, il_s:iu_s + ] # Restrict fine data else: @@ -833,17 +997,21 @@ def center_func_3(xm, xp): # Apply subsampling if subsample: # Calculate fine-level offsets (nearest cell at or below center) - o1 = s/2 - 1 if nx1 > 1 else 0 - o2 = s/2 - 1 if nx2 > 1 else 0 - o3 = s/2 - 1 if nx3 > 1 else 0 + o1 = s / 2 - 1 if nx1 > 1 else 0 + o2 = s / 2 - 1 if nx2 > 1 else 0 + o3 = s / 2 - 1 if nx3 > 1 else 0 # Assign values - for q, dataset, index in zip(quantities, quantity_datasets, - quantity_indices): - data[q][kl_d:ku_d, - jl_d:ju_d, - il_d:iu_d] = f[dataset][index, block_num, kl_s+o3:ku_s:s, - jl_s+o2:ju_s:s, il_s+o1:iu_s:s] + for q, dataset, index in zip( + quantities, quantity_datasets, quantity_indices + ): + data[q][kl_d:ku_d, jl_d:ju_d, il_d:iu_d] = f[dataset][ + index, + block_num, + kl_s + o3 : ku_s : s, + jl_s + o2 : ju_s : s, + il_s + o1 : iu_s : s, + ] # Apply fast (uniform Cartesian) restriction elif fast_restrict: @@ -853,18 +1021,22 @@ def center_func_3(xm, xp): ko_vals = list(range(s)) if nx3 > 1 else (0,) # Assign values - for q, dataset, index in zip(quantities, quantity_datasets, - quantity_indices): + for q, dataset, index in zip( + quantities, quantity_datasets, quantity_indices + ): for ko in ko_vals: for jo in jo_vals: for io in io_vals: - data[q][kl_d:ku_d, - jl_d:ju_d, - il_d:iu_d] += f[dataset][index, block_num, - kl_s+ko:ku_s:s, - jl_s+jo:ju_s:s, - il_s+io:iu_s:s] - data[q][kl_d:ku_d, jl_d:ju_d, il_d:iu_d] /= s ** num_extended_dims + data[q][kl_d:ku_d, jl_d:ju_d, il_d:iu_d] += f[ + dataset + ][ + index, + block_num, + kl_s + ko : ku_s : s, + jl_s + jo : ju_s : s, + il_s + io : iu_s : s, + ] + data[q][kl_d:ku_d, jl_d:ju_d, il_d:iu_d] /= s**num_extended_dims # Apply exact (volume-weighted) restriction else: @@ -885,24 +1057,24 @@ def center_func_3(xm, xp): # Accumulate values for k_s, k_d in zip(k_s_vals, k_d_vals): if nx3 > 1: - x3m = f['x3f'][block_num, k_s] - x3p = f['x3f'][block_num, k_s+1] + x3m = f["x3f"][block_num, k_s] + x3p = f["x3f"][block_num, k_s + 1] for j_s, j_d in zip(j_s_vals, j_d_vals): if nx2 > 1: - x2m = f['x2f'][block_num, j_s] - x2p = f['x2f'][block_num, j_s+1] + x2m = f["x2f"][block_num, j_s] + x2p = f["x2f"][block_num, j_s + 1] for i_s, i_d in zip(i_s_vals, i_d_vals): if nx1 > 1: - x1m = f['x1f'][block_num, i_s] - x1p = f['x1f'][block_num, i_s+1] + x1m = f["x1f"][block_num, i_s] + x1p = f["x1f"][block_num, i_s + 1] vol = vol_func(x1m, x1p, x2m, x2p, x3m, x3p) - for q, dataset, index in zip(quantities, - quantity_datasets, - quantity_indices): - data[q][k_d, j_d, i_d] += vol * f[dataset][index, - block_num, - k_s, j_s, - i_s] + for q, dataset, index in zip( + quantities, quantity_datasets, quantity_indices + ): + data[q][k_d, j_d, i_d] += ( + vol + * f[dataset][index, block_num, k_s, j_s, i_s] + ) loc1 = (nx1 > 1) * block_location[0] / s loc2 = (nx2 > 1) * block_location[1] / s loc3 = (nx3 > 1) * block_location[2] / s @@ -910,7 +1082,7 @@ def center_func_3(xm, xp): # Set level information for cells in this block if return_levels: - data['Levels'][kl_d:ku_d, jl_d:ju_d, il_d:iu_d] = block_level + data["Levels"][kl_d:ku_d, jl_d:ju_d, il_d:iu_d] = block_level # Remove volume factors from restricted data if level < max_level and not subsample and not fast_restrict: @@ -932,16 +1104,16 @@ def center_func_3(xm, xp): ku = min(ku, k_max) - k_min for k in range(kl, ku): if nx3 > 1: - x3m = data['x3f'][k] - x3p = data['x3f'][k+1] + x3m = data["x3f"][k] + x3p = data["x3f"][k + 1] for j in range(jl, ju): if nx2 > 1: - x2m = data['x2f'][j] - x2p = data['x2f'][j+1] + x2m = data["x2f"][j] + x2p = data["x2f"][j + 1] for i in range(il, iu): if nx1 > 1: - x1m = data['x1f'][i] - x1p = data['x1f'][i+1] + x1m = data["x1f"][i] + x1p = data["x1f"][i + 1] vol = vol_func(x1m, x1p, x2m, x2p, x3m, x3p) for q in quantities: data[q][k, j, i] /= vol @@ -957,6 +1129,7 @@ def center_func_3(xm, xp): # ======================================================================================== + def restrict_like(vals, levels, vols=None): """Average cell values according to given mesh refinement scheme.""" @@ -964,45 +1137,54 @@ def restrict_like(vals, levels, vols=None): nx3, nx2, nx1 = vals.shape max_level = np.max(levels) if nx3 > 1 and nx3 % 2**max_level != 0: - raise AthenaError('x3-dimension wrong size to be restricted') + raise AthenaError("x3-dimension wrong size to be restricted") if nx2 > 1 and nx2 % 2**max_level != 0: - raise AthenaError('x2-dimension wrong size to be restricted') + raise AthenaError("x2-dimension wrong size to be restricted") if nx1 % 2**max_level != 0: - raise AthenaError('x1-dimension wrong size to be restricted') + raise AthenaError("x1-dimension wrong size to be restricted") # Construct volume weighting if vols is None: vols = np.ones_like(vals) else: if vols.shape != vals.shape: - raise AthenaError('Array of volumes must match cell values in size') + raise AthenaError("Array of volumes must match cell values in size") # Restrict data vals_restricted = np.copy(vals) for level in range(max_level): level_difference = max_level - level - stride = 2 ** level_difference + stride = 2**level_difference if nx3 > 1: - vals_level = np.reshape(vals * vols, (nx3/stride, stride, nx2/stride, stride, - nx1/stride, stride)) - vols_level = np.reshape(vols, (nx3/stride, stride, nx2/stride, stride, - nx1/stride, stride)) + vals_level = np.reshape( + vals * vols, + (nx3 / stride, stride, nx2 / stride, stride, nx1 / stride, stride), + ) + vols_level = np.reshape( + vols, (nx3 / stride, stride, nx2 / stride, stride, nx1 / stride, stride) + ) vals_sum = np.sum(np.sum(np.sum(vals_level, axis=5), axis=3), axis=1) vols_sum = np.sum(np.sum(np.sum(vols_level, axis=5), axis=3), axis=1) - vals_level = np.repeat(np.repeat(np.repeat(vals_sum / vols_sum, stride, - axis=0), - stride, axis=1), - stride, axis=2) + vals_level = np.repeat( + np.repeat( + np.repeat(vals_sum / vols_sum, stride, axis=0), stride, axis=1 + ), + stride, + axis=2, + ) elif nx2 > 1: - vals_level = np.reshape(vals * vols, (nx2/stride, stride, nx1/stride, stride)) - vols_level = np.reshape(vols, (nx2/stride, stride, nx1/stride, stride)) + vals_level = np.reshape( + vals * vols, (nx2 / stride, stride, nx1 / stride, stride) + ) + vols_level = np.reshape(vols, (nx2 / stride, stride, nx1 / stride, stride)) vals_sum = np.sum(np.sum(vals_level, axis=3), axis=1) vols_sum = np.sum(np.sum(vols_level, axis=3), axis=1) - vals_level = np.repeat(np.repeat(vals_sum / vols_sum, stride, axis=0), - stride, axis=1) + vals_level = np.repeat( + np.repeat(vals_sum / vols_sum, stride, axis=0), stride, axis=1 + ) else: - vals_level = np.reshape(vals * vols, (nx1/stride, stride)) - vols_level = np.reshape(vols, (nx1/stride, stride)) + vals_level = np.reshape(vals * vols, (nx1 / stride, stride)) + vols_level = np.reshape(vols, (nx1 / stride, stride)) vals_sum = np.sum(vals_level, axis=1) vols_sum = np.sum(vols_level, axis=1) vals_level = np.repeat(vals_sum / vols_sum, stride, axis=0) @@ -1012,21 +1194,23 @@ def restrict_like(vals, levels, vols=None): # ======================================================================================== + def athinput(filename): """Read athinput file and returns a dictionary of dictionaries.""" # Read data - with open(filename, 'r') as athinput: + with open(filename, "r") as athinput: # remove comments, extra whitespace, and empty lines - lines = [_f for _f in - [i.split('#')[0].strip() for i in athinput.readlines()] if _f] + lines = [ + _f for _f in [i.split("#")[0].strip() for i in athinput.readlines()] if _f + ] data = {} # split into blocks, first element will be empty - blocks = ('\n'.join(lines)).split('<')[1:] + blocks = ("\n".join(lines)).split("<")[1:] # Function for interpreting strings numerically def typecast(x): - if '_' in x: + if "_" in x: return x try: return int(x) @@ -1044,14 +1228,14 @@ def typecast(x): # Function for parsing assignment based on first '=' def parse_line(line): - out = [i.strip() for i in line.split('=')] - out[1] = '='.join(out[1:]) + out = [i.strip() for i in line.split("=")] + out[1] = "=".join(out[1:]) out[1] = typecast(out[1]) return out[:2] # Assign values into dictionaries for block in blocks: - info = list([_f for _f in block.split('\n') if _f]) + info = list([_f for _f in block.split("\n") if _f]) key = info.pop(0)[:-1] # last character is '>' data[key] = dict(list(map(parse_line, info))) return data @@ -1059,11 +1243,14 @@ def parse_line(line): # ======================================================================================== + class AthenaError(RuntimeError): """General exception class for Athena++ read functions.""" + pass class AthenaWarning(RuntimeWarning): """General warning class for Athena++ read functions.""" + pass diff --git a/astro_tigress/map_circular_beam.py b/astro_tigress/map_circular_beam.py index 186d133..a90d72f 100644 --- a/astro_tigress/map_circular_beam.py +++ b/astro_tigress/map_circular_beam.py @@ -2,57 +2,60 @@ import math from scipy.integrate import quad + def get_gauss_stamp(n): nx = n ny = n - nxc = nx/2 - nyc = ny/2 + nxc = nx / 2 + nyc = ny / 2 ix = np.zeros((nx, ny)) iy = np.zeros((nx, ny)) for i in range(nx): for j in range(ny): - ix[i, j] = j-nyc - iy[i, j] = i-nxc + ix[i, j] = j - nyc + iy[i, j] = i - nxc + def gauss(x): sigmax = 0.5 - ret = 1./(np.sqrt(2*math.pi)*sigmax) * np.exp(-0.5*(x/sigmax)**2) + ret = 1.0 / (np.sqrt(2 * math.pi) * sigmax) * np.exp(-0.5 * (x / sigmax) ** 2) return ret + stamp = np.zeros((nx, ny)) for i in range(nx): for j in range(ny): x1 = ix[i, j] - 0.5 - x2 = x1 + 1. + x2 = x1 + 1.0 y1 = iy[i, j] - 0.5 - y2 = y1 + 1. + y2 = y1 + 1.0 inte = quad(gauss, x1, x2)[0] * quad(gauss, y1, y2)[0] stamp[i, j] = inte return stamp -#average data using a stamp + +# average data using a stamp def stamp_avg(data, stamp): nxd, nyd = data.shape nxs, nys = stamp.shape ret = np.zeros(data.shape) - nxc = nxs/2 - nyc = nys/2 + nxc = nxs / 2 + nyc = nys / 2 ix = np.zeros((nxs, nys)) iy = np.zeros((nxs, nys)) for i in range(nxs): for j in range(nys): - ix[i, j] = i-nxc - iy[i, j] = j-nyc + ix[i, j] = i - nxc + iy[i, j] = j - nyc for i in range(nxd): for j in range(nyd): for istamp in range(nxs): for jstamp in range(nys): iret = i + ix[istamp, jstamp] jret = j + iy[istamp, jstamp] - if (iret >= 0) and (iret < nxd - ) and (jret >= 0) and (jret < nyd): - ret[iret, jret] += data[i, j]*stamp[istamp, jstamp] + if (iret >= 0) and (iret < nxd) and (jret >= 0) and (jret < nyd): + ret[iret, jret] += data[i, j] * stamp[istamp, jstamp] return ret + def map_circular_beam(data, nstamp=9): stamp = get_gauss_stamp(nstamp) return stamp_avg(data, stamp) - diff --git a/astro_tigress/tigress_read.py b/astro_tigress/tigress_read.py index f963937..c573375 100644 --- a/astro_tigress/tigress_read.py +++ b/astro_tigress/tigress_read.py @@ -168,9 +168,7 @@ def _download_file(self, f): target = osp.join(self.dir_master, f) if osp.isfile(target): - print( - "{} ({:.5f}GB) already exists".format(f, osp.getsize(target) / 2**30) - ) + print("{} ({:.5f}GB) already exists".format(f, osp.getsize(target) / 2**30)) while True: answer = input("overwrite? [y/n]:") if answer.lower() in ["y", "n"]: @@ -184,7 +182,7 @@ def _download_file(self, f): req = urllib.request.Request(source) try: - response = urllib.request.urlopen(req) + urllib.request.urlopen(req) except URLError as e: if hasattr(e, "reason"): print("We failed to reach a server.") diff --git a/astro_tigress/tigress_utils.py b/astro_tigress/tigress_utils.py index d3da4c0..c985b49 100644 --- a/astro_tigress/tigress_utils.py +++ b/astro_tigress/tigress_utils.py @@ -1,6 +1,7 @@ import os import shutil + def copy_file(fn_old, fn_new): """Copy file from old to new location. Make new directory if needed. Raise warning if old file does not exist. @@ -8,20 +9,20 @@ def copy_file(fn_old, fn_new): fn_old: string, old filename fn_new: string, new filename""" if fn_old == fn_new: - raise ValueError('New filename cannot be the same as the original one.') + raise ValueError("New filename cannot be the same as the original one.") if os.path.isfile(fn_old): if os.path.isfile(fn_new): - print("New file already exists: "+fn_new) + print("New file already exists: " + fn_new) while True: answer = input("Overwrite? (y/n):") if answer == "y": break elif answer == "n": - return + return dir_new = os.path.dirname(fn_new) if not os.path.exists(dir_new): os.makedirs(dir_new) shutil.copyfile(fn_old, fn_new) else: - raise RuntimeError("file does not exist: "+fn_old) + raise RuntimeError("file does not exist: " + fn_old) return diff --git a/astro_tigress/ytathena.py b/astro_tigress/ytathena.py index f44f076..75d488a 100644 --- a/astro_tigress/ytathena.py +++ b/astro_tigress/ytathena.py @@ -39,7 +39,7 @@ def __init__(self, fname=os.path.join(dirpath, "coolftn.txt")): self.nT = len(self.T1) def get_Tidx(self, T): - if type(T) == np.ndarray: + if type(T) is np.ndarray: Tidx = np.log10(T / self.Tmin) / self.dT Tidx[np.where(T < self.Tmin)] = 0 Tidx[np.where(T >= self.Tmax)] = self.nT - 2