def sphere_sum(dump, var, r_slice=None, i_slice=None, th_slice=None, j_slice=None, mask=None): """Sum everything within a sphere, semi-sphere, or thick spherical shell Extent can be specified in r and/or theta, or i and/or j Mask is multiplied at the end """ # TODO see sum for problems with this if th_slice is not None: j_slice = get_eht_disk_j_vals(dump, th_slice[0], th_slice[1]) if j_slice is not None: var = var[:, j_slice[0]:j_slice[1], :] gdet = dump['gdet'][:, j_slice[0]:j_slice[1]] else: gdet = dump['gdet'] if r_slice is not None: i_slice = (i_of(dump['r'][:, 0, 0], r_slice[0]), i_of(dump['r'][:, 0, 0], r_slice[1])) if i_slice is not None: var = var[i_slice[0]:i_slice[1]] gdet = gdet[i_slice[0]:i_slice[1]] return np.sum(var * gdet[:, :, None] * dump.header['dx1'] * dump.header['dx2'] * dump.header['dx3'])
def get_eht_disk_j_vals(dump, th_min=np.pi / 3., th_max=2 * np.pi / 3.): """Calculate jmin, jmax in theta coordinate for EHT disk profiles (pi/3 - 2pi/3)""" # Calculate jmin, jmax for EHT radial profiles # Use values at large R to get even split in radially-dependent coordinates if len(dump['th'].shape) == 3: ths = dump['th'][-1, :, 0] elif len(dump['th'].shape) == 2: ths = dump['th'][-1, :] return (i_of(ths, th_min), i_of(ths, th_max))
def get_quiescence(infname, diag=False, set_time=None): with h5py.File(infname, 'r') as infile: if set_time is not None: tstart, tend = set_time else: tstart, tend = infile['avg']['start'][()], infile['avg']['end'][()] if diag: t = infile['diag']['t'][()] else: t = infile['coord']['t'][()] start = i_of(t, tstart) end = i_of(t, tend) return slice(start, end)
def shell_sum(dump, var, at_r=None, at_zone=None, th_slice=None, j_slice=None, mask=None): """Sum a variable over spherical shells. Returns a radial profile (array length N1) or single-shell sum @param at_r: Single radius at which to sum (nearest-neighbor smaller zone is used) @param at_zone: Specific radial zone at which to sum, for compatibility @param th_slice: Tuple of minimum and maximum theta value to sum @param j_slice: Tuple of x2 indices instead of specifying theta @param mask: array of 1/0 of remaining size which is multiplied with the result """ if isinstance(var, str): var = dump[var] # Translate coordinates to zone numbers. # TODO Xtoijk for slices? # TODO slice dx2, dx3 if they're matrices for exotic coordinates if th_slice is not None: j_slice = get_eht_disk_j_vals(dump, th_slice[0], th_slice[1]) if j_slice is not None: var = var[:, j_slice[0]:j_slice[1], :] gdet = dump['gdet'][:, j_slice[0]:j_slice[1]] else: gdet = dump['gdet'] if at_r is not None: at_zone = i_of(dump['r'][:, 0, 0], at_r) if at_zone is not None: # Keep integrand "3D" and deal with it below var = var[at_zone:at_zone + 1] gdet = gdet[at_zone:at_zone + 1] integrand = var * gdet[:, :, None] * dump.header['dx2'] * dump.header['dx3'] if mask is not None: integrand *= mask ret = np.sum(integrand, axis=(-2, -1)) if ret.shape == (1, ): # Don't return a scalar result as a length-1 array return ret[0] else: return ret
def get_ivar(infname, ivar, th_r=None, i_xy=False, mesh=True): """Given an input file and the string of independent variable name(s) ('r', 'rth', 'rt', etc), return a grid of those variables' values. """ ret_i = [] G = get_grid(infname) if ivar[:4] == "log_": do_log = True else: do_log = False if mesh: native_coords = G.coord_all_mesh() else: native_coords = G.coord_all() if ivar[-1:] == 't': with h5py.File(infname, 'r') as infile: t = infile['coord']['t'][()] if mesh: t = np.append(t, t[-1] + (t[-1] - t[0]) / t.shape[0]) ret_i.append(t) if 'r' in ivar: ret_i.append(G.coords.r(native_coords)[:, 0, 0]) if 'th' in ivar: r1d = G.coords.r(native_coords)[:, 0, 0] if th_r is not None: th = G.coords.th(native_coords)[i_of(r1d, th_r), :, 0] else: #print("Guessing r for computing th!") th = G.coords.th(native_coords)[-1, :, 0] if 'hth' in ivar: th = th[:len(th) // 2] ret_i.append(th) if 'phi' in ivar: ret_i.append(G.coords.phi(native_coords)[0, 0, :]) # TODO handle converting 'thphi' to x-y with at_r # TODO handle th's r-dependence in 'rth' # TODO think about how to treat slices this nicely # Make a meshgrid of ret_grids = np.meshgrid(*reversed(ret_i)) ret_grids.reverse() if i_xy and 'r' in ivar and 'th' in ivar: # XZ plot x = ret_grids[-2] * np.sin(ret_grids[-1]) z = ret_grids[-2] * np.cos(ret_grids[-1]) ret_grids[-2:] = x, z elif i_xy and 'r' in ivar and 'phi' in ivar: # XY plot x = ret_grids[-2] * np.cos(ret_grids[-1]) y = ret_grids[-2] * np.sin(ret_grids[-1]) ret_grids[-2:] = x, y if do_log: #print("Applying logs shape", len(ret_grids)) for i in range(len(ret_grids)): ret_grids[i] = np.log(ret_grids[i]) # Squash single-variable lists for convenience if len(ret_grids) == 1: ret_grids = ret_grids[0] if mesh: # Probably no one actually wants a 1D mesh ret_grids = ret_grids[:-1] return ret_grids
def plot(n): tdump = io.get_dump_time(files[n]) if (tstart is not None and tdump < tstart) or (tend is not None and tdump > tend): return print("frame {} / {}".format(n, len(files) - 1)) fig = plt.figure(figsize=(FIGX, FIGY)) to_load = {} if "simple" not in movie_type and "floor" not in movie_type: # Everything but simple & pure floor movies needs derived vars to_load['calc_derived'] = True if "simple" in movie_type: # Save memory #to_load['add_grid_caches'] = False pass if "fail" in movie_type or "e_ratio" in movie_type or "conservation" in movie_type: to_load['add_fails'] = True if "floor" in movie_type: to_load['add_floors'] = True if "current" in movie_type or "jsq" in movie_type or "jcon" in movie_type: to_load['add_jcon'] = True if "divB" in movie_type: to_load['add_divB'] = True #to_load['calc_divB'] = True if "psi_cd" in movie_type: to_load['add_psi_cd'] = True if "1d" in movie_type: to_load['add_grid_caches'] = False to_load['calc_derived'] = False if "_ghost" in movie_type: plot_ghost = True to_load['add_ghosts'] = True else: plot_ghost = False # TODO U if needed dump = pyHARM.load_dump(files[n], **to_load) # Title by time, otherwise number #try: # fig.suptitle("t = {}".format(int(dump['t']))) #except ValueError: # fig.suptitle("dump {}".format(n)) # Zoom in for small problems # TODO use same r1d as analysis? if len(dump['r'].shape) == 1: r1d = dump['r'] sz = 50 nlines = 20 rho_l, rho_h = None, None elif len(dump['r'].shape) == 2: r1d = dump['r'][:, 0] sz = 50 nlines = 20 rho_l, rho_h = -6, 1 else: r1d = dump['r'][:, 0, 0] if dump['r'][-1, 0, 0] > 100: sz = 50 nlines = 20 rho_l, rho_h = -5, 1.5 iBZ = i_of(r1d, 100) # most MADs rBZ = 100 elif dump['r'][-1, 0, 0] > 10: sz = 50 nlines = 5 rho_l, rho_h = -6, 1 iBZ = i_of(r1d, 40) # most SANEs rBZ = 40 else: # Then this is a Minkowski simulation or something weird. Guess. sz = (dump['x'][-1, 0, 0] - dump['x'][0, 0, 0]) / 2 nlines = 0 rho_l, rho_h = -2, 0.0 iBZ = 1 rBZ = 1 window = [-sz, sz, -sz, sz] # If we're in arrspace we (almost) definitely want a 0,1 window # TODO allow zooming in toward corners. Original r vs th as separate plotting set? if "_array" in movie_type: USEARRSPACE = True if plot_ghost: window = [-0.1, 1.1, -0.1, 1.1] else: window = [0, 1, 0, 1] else: USEARRSPACE = False if movie_type == "simplest_poloidal": # Simplest movie: just RHO, poloidal slice ax_slc = plt.subplot(1, 1, 1) var = 'rho' arrspace = False vmin = None vmax = None pplt.plot_xz(ax_slc, dump, var, label="", vmin=vmin, vmax=vmax, window=window, arrayspace=arrspace, xlabel=False, ylabel=False, xticks=[], yticks=[], cbar=False, cmap='jet', use_imshow=True) ax_slc.axis('off') plt.subplots_adjust(hspace=0, wspace=0, left=0, right=1, bottom=0, top=1) elif movie_type == "simplest_toroidal": # Simplest movie: just RHO, toroidal slice ax_slc = plt.subplot(1, 1, 1) var = 'log_rho' arrspace = False vmin = rho_l vmax = rho_h pplt.plot_xy(ax_slc, dump, var, label="", vmin=vmin + 0.15, vmax=vmax + 0.15, window=window, arrayspace=arrspace, xlabel=False, ylabel=False, xticks=[], yticks=[], cbar=False, cmap='jet') plt.subplots_adjust(hspace=0, wspace=0, left=0, right=1, bottom=0, top=1) elif movie_type == "simplest": # Simplest movie: just RHO ax_slc = [plt.subplot(1, 2, 1), plt.subplot(1, 2, 2)] if dump['coordinates'] == "cartesian": var = 'rho' arrspace = True vmin = None vmax = None else: arrspace = USEARRSPACE # Linear version # var = 'rho' # vmin = 0 # vmax = 1 var = 'log_rho' vmin = rho_l vmax = rho_h pplt.plot_xz(ax_slc[0], dump, var, label="", vmin=vmin, vmax=vmax, window=window, arrayspace=arrspace, xlabel=False, ylabel=False, xticks=[], yticks=[], cbar=False, cmap='jet') pplt.plot_xy(ax_slc[1], dump, var, label="", vmin=vmin + 0.15, vmax=vmax + 0.15, window=window, arrayspace=arrspace, xlabel=False, ylabel=False, xticks=[], yticks=[], cbar=False, cmap='jet') pad = 0.0 plt.subplots_adjust(hspace=0, wspace=0, left=pad, right=1 - pad, bottom=pad, top=1 - pad) elif movie_type == "simpler": # Simpler movie: RHO and phi gs = gridspec.GridSpec(2, 2, height_ratios=[6, 1], width_ratios=[16, 17]) ax_slc = [fig.subplot(gs[0, 0]), fig.subplot(gs[0, 1])] ax_flux = [fig.subplot(gs[1, :])] pplt.plot_slices(ax_slc[0], ax_slc[1], dump, 'log_rho', vmin=rho_l, vmax=rho_h, window=window, overlay_field=False, cmap='jet') ppltr.plot_diag(ax_flux[0], diag, 'phi_b', tline=dump['t'], logy=LOG_PHI, xlabel=False) elif movie_type == "simple": # Simple movie: RHO mdot phi gs = gridspec.GridSpec(3, 2, height_ratios=[4, 1, 1]) ax_slc = [fig.subplot(gs[0, 0]), fig.subplot(gs[0, 1])] ax_flux = [fig.subplot(gs[1, :]), fig.subplot(gs[2, :])] pplt.plot_slices(ax_slc[0], ax_slc[1], dump, 'log_rho', vmin=rho_l, vmax=rho_h, window=window, cmap='jet', arrayspace=USEARRSPACE) ppltr.plot_diag(ax_flux[0], diag, 'Mdot', tline=dump['t'], logy=LOG_MDOT) ppltr.plot_diag(ax_flux[1], diag, 'Phi_b', tline=dump['t'], logy=LOG_PHI) elif movie_type == "traditional" or movie_type == "eht": ax_slc = lambda i: plt.subplot(2, 4, i) # Usual movie: RHO beta fluxes # CUTS pplt.plot_slices(ax_slc(1), ax_slc(2), dump, 'log_rho', label='log_rho', average=True, vmin=rho_l, vmax=rho_h, cmap='jet', window=window, arrayspace=USEARRSPACE) pplt.plot_slices(ax_slc(3), ax_slc(4), dump, 'log_UU', label='log_UU', average=True, vmin=rho_l, vmax=rho_h, cmap='jet', window=window, arrayspace=USEARRSPACE) pplt.plot_slices(ax_slc(5), ax_slc(6), dump, 'log_bsq', label='log_bsq', average=True, vmin=rho_l, vmax=rho_h, cmap='jet', window=window, arrayspace=USEARRSPACE) pplt.plot_slices(ax_slc(7), ax_slc(8), dump, 'log_beta', label='log_beta', average=True, vmin=rho_l, vmax=rho_h, cmap='jet', window=window, arrayspace=USEARRSPACE) # FLUXES # ppltr.plot_diag(ax_flux(2), diag, 't', 'Mdot', tline=dump['t'], logy=LOG_MDOT) # ppltr.plot_diag(ax_flux(4), diag, 't', 'phi_b', tline=dump['t'], logy=LOG_PHI) # Mixins: # Zoomed in RHO # pplt.plot_slices(ax_slc(7), ax_slc(8), dump, 'log_rho', vmin=-3, vmax=2, # window=[-10, 10, -10, 10], field_overlay=False) elif movie_type == "prims_xz": ax_slc = lambda i: plt.subplot(2, 4, i) vmin, vmax = None, None pplt.plot_xz(ax_slc(1), dump, 'RHO', label="", vmin=vmin, vmax=vmax, window=window, arrayspace=USEARRSPACE, xlabel=False, ylabel=False, xticks=[], yticks=[], cbar=False, cmap='jet') pplt.plot_xz(ax_slc(2), dump, 'UU', label="", vmin=vmin, vmax=vmax, window=window, arrayspace=USEARRSPACE, xlabel=False, ylabel=False, xticks=[], yticks=[], cbar=False, cmap='jet') pplt.plot_xz(ax_slc(3), dump, 'U1', label="", vmin=vmin, vmax=vmax, window=window, arrayspace=USEARRSPACE, xlabel=False, ylabel=False, xticks=[], yticks=[], cbar=False, cmap='jet') pplt.plot_xz(ax_slc(4), dump, 'U2', label="", vmin=vmin, vmax=vmax, window=window, arrayspace=USEARRSPACE, xlabel=False, ylabel=False, xticks=[], yticks=[], cbar=False, cmap='jet') pplt.plot_xz(ax_slc(5), dump, 'U3', label="", vmin=vmin, vmax=vmax, window=window, arrayspace=USEARRSPACE, xlabel=False, ylabel=False, xticks=[], yticks=[], cbar=False, cmap='jet') pplt.plot_xz(ax_slc(6), dump, 'B1', label="", vmin=vmin, vmax=vmax, window=window, arrayspace=USEARRSPACE, xlabel=False, ylabel=False, xticks=[], yticks=[], cbar=False, cmap='jet') pplt.plot_xz(ax_slc(7), dump, 'B2', label="", vmin=vmin, vmax=vmax, window=window, arrayspace=USEARRSPACE, xlabel=False, ylabel=False, xticks=[], yticks=[], cbar=False, cmap='jet') pplt.plot_xz(ax_slc(8), dump, 'B3', label="", vmin=vmin, vmax=vmax, window=window, arrayspace=USEARRSPACE, xlabel=False, ylabel=False, xticks=[], yticks=[], cbar=False, cmap='jet') elif movie_type == "prims_xz_array": ax_slc = lambda i: plt.subplot(2, 4, i) vmin, vmax = None, None pplt.plot_xz(ax_slc(1), dump, 'RHO', label="", vmin=vmin, vmax=vmax, window=window, arrayspace=True, xlabel=False, ylabel=False, xticks=[], yticks=[], cbar=False, cmap='jet') pplt.plot_xz(ax_slc(2), dump, 'UU', label="", vmin=vmin, vmax=vmax, window=window, arrayspace=True, xlabel=False, ylabel=False, xticks=[], yticks=[], cbar=False, cmap='jet') pplt.plot_xz(ax_slc(3), dump, 'U1', label="", vmin=vmin, vmax=vmax, window=window, arrayspace=True, xlabel=False, ylabel=False, xticks=[], yticks=[], cbar=False, cmap='jet') pplt.plot_xz(ax_slc(4), dump, 'U2', label="", vmin=vmin, vmax=vmax, window=window, arrayspace=True, xlabel=False, ylabel=False, xticks=[], yticks=[], cbar=False, cmap='jet') pplt.plot_xz(ax_slc(5), dump, 'U3', label="", vmin=vmin, vmax=vmax, window=window, arrayspace=True, xlabel=False, ylabel=False, xticks=[], yticks=[], cbar=False, cmap='jet') pplt.plot_xz(ax_slc(6), dump, 'B1', label="", vmin=vmin, vmax=vmax, window=window, arrayspace=True, xlabel=False, ylabel=False, xticks=[], yticks=[], cbar=False, cmap='jet') pplt.plot_xz(ax_slc(7), dump, 'B2', label="", vmin=vmin, vmax=vmax, window=window, arrayspace=True, xlabel=False, ylabel=False, xticks=[], yticks=[], cbar=False, cmap='jet') pplt.plot_xz(ax_slc(8), dump, 'B3', label="", vmin=vmin, vmax=vmax, window=window, arrayspace=True, xlabel=False, ylabel=False, xticks=[], yticks=[], cbar=False, cmap='jet') elif movie_type == "vectors": ax_slc = lambda i: plt.subplot(2, 4, i) # Usual movie: RHO beta fluxes # CUTS pplt.plot_slices(ax_slc(1), ax_slc(5), dump, 'log_rho', label=pretty('log_rho'), average=True, vmin=rho_l, vmax=rho_h, cmap='jet', window=window, arrayspace=USEARRSPACE) for i, var in zip((2, 3, 4, 6, 7, 8), ("U1", "U2", "U3", "B1", "B2", "B3")): pplt.plot_xz(ax_slc(i), dump, np.log10(dump[var]), label=var, vmin=rho_l, vmax=rho_h, cmap='Reds', window=window, arrayspace=USEARRSPACE) pplt.plot_xz(ax_slc(i), dump, np.log10(-dump[var]), label=var, vmin=rho_l, vmax=rho_h, cmap='Blues', window=window, arrayspace=USEARRSPACE) elif movie_type == "vecs_cov": ax_slc = lambda i: plt.subplot(2, 4, i) for i, var in zip( (1, 2, 3, 4, 5, 6, 7, 8), ("u_0", "u_r", "u_th", "u_3", "b_0", "b_r", "b_th", "b_3")): pplt.plot_xz(ax_slc(i), dump, np.log10(dump[var]), label=var, vmin=rho_l, vmax=rho_h, cmap='Reds', window=window, arrayspace=USEARRSPACE) pplt.plot_xz(ax_slc(i), dump, np.log10(-dump[var]), label=var, vmin=rho_l, vmax=rho_h, cmap='Blues', window=window, arrayspace=USEARRSPACE) elif movie_type == "vecs_con": ax_slc = lambda i: plt.subplot(2, 4, i) for i, var in zip( (1, 2, 3, 4, 5, 6, 7, 8), ("u^0", "u^r", "u^th", "u^3", "b^0", "b^r", "b^th", "b^3")): pplt.plot_xz(ax_slc(i), dump, np.log10(dump[var]), label=var, vmin=rho_l, vmax=rho_h, cmap='Reds', window=window, arrayspace=USEARRSPACE) pplt.plot_xz(ax_slc(i), dump, np.log10(-dump[var]), label=var, vmin=rho_l, vmax=rho_h, cmap='Blues', window=window, arrayspace=USEARRSPACE) elif movie_type == "ejection": ax_slc = lambda i: plt.subplot(1, 2, i) # Usual movie: RHO beta fluxes # CUTS pplt.plot_xz(ax_slc(1), dump, 'log_rho', label=pretty('log_rho') + " phi-average", average=True, vmin=rho_l, vmax=rho_h, cmap='jet', window=window, arrayspace=USEARRSPACE) pplt.plot_xz(ax_slc(2), dump, 'log_bsq', label=pretty('log_bsq') + " phi-average", average=True, vmin=rho_l, vmax=rho_h, cmap='jet', window=window, arrayspace=USEARRSPACE) elif movie_type == "b_bug": rmax = 10 thmax = 10 phi = 100 ax_slc = lambda i: plt.subplot(1, 3, i) ax_slc(1).pcolormesh(dump['X1'][:rmax, 0:thmax, phi], dump['X2'][:rmax, 0:thmax, phi], dump['log_b^r'][:rmax, 0:thmax, phi], vmax=0, vmin=-4) ax_slc(2).pcolormesh(dump['X1'][:rmax, 0:thmax, phi], dump['X2'][:rmax, 0:thmax, phi], dump['log_b^th'][:rmax, 0:thmax, phi], vmax=0, vmin=-4) ax_slc(3).pcolormesh(dump['X1'][:rmax, 0:thmax, phi], dump['X2'][:rmax, 0:thmax, phi], dump['log_b^3'][:rmax, 0:thmax, phi], vmax=0, vmin=-4) elif movie_type == "e_ratio": ax_slc = lambda i: plt.subplot(2, 4, i) # Energy ratios: difficult places to integrate, with failures pplt.plot_slices(ax_slc(1), ax_slc(2), dump, np.log10(dump['UU'] / dump['RHO']), label=r"$\log_{10}(U / \rho)$", vmin=-3, vmax=3, average=True, field_overlay=False, window=window, arrayspace=USEARRSPACE) pplt.plot_slices(ax_slc(3), ax_slc(4), dump, np.log10(dump['bsq'] / dump['RHO']), label=r"$\log_{10}(b^2 / \rho)$", vmin=-3, vmax=3, average=True, field_overlay=False, window=window, arrayspace=USEARRSPACE) pplt.plot_slices(ax_slc(5), ax_slc(6), dump, np.log10(1 / dump['beta']), label=r"$\beta^{-1}$", vmin=-3, vmax=3, average=True, field_overlay=False, window=window, arrayspace=USEARRSPACE) pplt.plot_slices(ax_slc(7), ax_slc(8), dump, (dump['fails'] != 0).astype(np.int32), label="Failures", vmin=0, vmax=20, cmap='Reds', integrate=True, field_overlay=False, window=window, arrayspace=USEARRSPACE) elif "e_ratio_funnel" in movie_type: ax_slc = lambda i: plt.subplot(2, 4, i) # Energy ratios: difficult places to integrate, with failures r_i = i_of(r1d, float(movie_type.split("_")[-1])) pplt.plot_thphi(ax_slc(1), dump, np.log10(dump['UU'] / dump['RHO']), r_i, label=r"$\log_{10}(U / \rho)$", vmin=-3, vmax=3, average=True, field_overlay=False, window=window, arrayspace=USEARRSPACE) pplt.plot_thphi(ax_slc(2), dump, np.log10(dump['UU'] / dump['RHO']), r_i, label=r"$\log_{10}(U / \rho)$", vmin=-3, vmax=3, average=True, field_overlay=False, window=window, arrayspace=USEARRSPACE) pplt.plot_thphi(ax_slc(3), dump, np.log10(dump['bsq'] / dump['RHO']), r_i, label=r"$\log_{10}(b^2 / \rho)$", vmin=-3, vmax=3, average=True, field_overlay=False, window=window, arrayspace=USEARRSPACE) pplt.plot_thphi(ax_slc(4), dump, np.log10(dump['bsq'] / dump['RHO']), r_i, label=r"$\log_{10}(b^2 / \rho)$", vmin=-3, vmax=3, average=True, field_overlay=False, window=window, arrayspace=USEARRSPACE) pplt.plot_thphi(ax_slc(5), dump, np.log10(1 / dump['beta']), r_i, label=r"$\beta^{-1}$", vmin=-3, vmax=3, average=True, field_overlay=False, window=window, arrayspace=USEARRSPACE) pplt.plot_thphi(ax_slc(6), dump, np.log10(1 / dump['beta']), r_i, label=r"$\beta^{-1}$", vmin=-3, vmax=3, average=True, field_overlay=False, window=window, arrayspace=USEARRSPACE) pplt.plot_thphi(ax_slc(7), dump, (dump['fails'] != 0).astype(np.int32), r_i, label="Failures", vmin=0, vmax=20, cmap='Reds', integrate=True, field_overlay=False, window=window, arrayspace=USEARRSPACE) pplt.plot_thphi(ax_slc(8), dump, (dump['fails'] != 0).astype(np.int32), r_i, label="Failures", vmin=0, vmax=20, cmap='Reds', integrate=True, field_overlay=False, window=window, arrayspace=USEARRSPACE) elif movie_type == "conservation": ax_slc = lambda i: plt.subplot(2, 4, i) ax_flux = lambda i: plt.subplot(4, 2, i) # Continuity plots to verify local conservation of energy, angular + linear momentum # Integrated T01: continuity for momentum conservation pplt.plot_slices(ax_slc(1), ax_slc(2), dump, T_mixed(dump, 1, 0), label=r"$T^1_0$ Integrated", vmin=0, vmax=2000, arrspace=True, integrate=True) # integrated T00: continuity plot for energy conservation pplt.plot_slices(ax_slc(5), ax_slc(6), dump, np.abs(T_mixed(dump, 0, 0)), label=r"$T^0_0$ Integrated", vmin=0, vmax=3000, arrspace=True, integrate=True) # Usual fluxes for reference #ppltr.plot_diag(ax_flux[1], diag, 't', 'mdot', tline=dump['t'], logy=LOG_MDOT) r_out = 100 # Radial conservation plots E_r = shell_sum(dump, T_mixed(dump, 0, 0)) # TODO variables Ang_r = shell_sum(dump, T_mixed(dump, 0, 3)) mass_r = shell_sum(dump, dump['ucon'][0] * dump['RHO']) max_e = 50000 pplt.radial_plot(ax_flux(2), dump, np.abs(E_r), title='Conserved vars at R', ylim=(0, max_e), rlim=(0, r_out), label="E_r") pplt.radial_plot(ax_flux(2), dump, np.abs(Ang_r) / 10, ylim=(0, max_e), rlim=(0, r_out), color='r', label="L_r") pplt.radial_plot(ax_flux(2), dump, np.abs(mass_r), ylim=(0, max_e), rlim=(0, r_out), color='b', label="M_r") ax_flux(2).legend() # Radial energy accretion rate Edot_r = shell_sum(dump, T_mixed(dump, 1, 0)) pplt.radial_plot(ax_flux(4), dump, Edot_r, label='Edot at R', ylim=(-200, 200), rlim=(0, r_out), arrayspace=True) # Radial integrated failures pplt.radial_plot(ax_flux(6), dump, (dump['fails'] != 0).sum(axis=(1, 2)), label='Fails at R', arrayspace=True, rlim=(0, r_out), ylim=(0, 1000)) elif movie_type == "energies": ax_slc = lambda i: plt.subplot(2, 4, i) # Energy ratios: difficult places to integrate, with failures pplt.plot_slices(ax_slc(1), ax_slc(2), dump, 'log_rho', label=r"$\log_{10}(U / \rho)$", vmin=-3, vmax=3, average=True, field_overlay=False, window=window, arrayspace=USEARRSPACE) pplt.plot_slices(ax_slc(3), ax_slc(4), dump, 'log_bsq', label=r"$\log_{10}(b^2 / \rho)$", vmin=-3, vmax=3, average=True, field_overlay=False, window=window, arrayspace=USEARRSPACE) pplt.plot_slices(ax_slc(5), ax_slc(6), dump, 'log_UU', label=r"$\beta^{-1}$", vmin=-3, vmax=3, average=True, field_overlay=False, window=window, arrayspace=USEARRSPACE) pplt.plot_slices(ax_slc(7), ax_slc(8), dump, (dump['fails'] != 0).astype(np.int32), label="Failures", vmin=0, vmax=20, cmap='Reds', integrate=True, field_overlay=False, window=window, arrayspace=USEARRSPACE) elif movie_type == "floors": ax_slc = lambda i: plt.subplot(2, 4, i) pplt.plot_xz(ax_slc(1), dump, 'log_rho', label=pretty('log_rho'), vmin=rho_l, vmax=rho_h, cmap='jet', window=window, arrayspace=USEARRSPACE) max_fail = 20 pplt.plot_xz(ax_slc(2), dump, dump['floors'] & 1, label="GEOM_RHO", vmin=0, vmax=max_fail, cmap='Reds', integrate=True, window=window, arrayspace=USEARRSPACE) pplt.plot_xz(ax_slc(3), dump, dump['floors'] & 2, label="GEOM_U", vmin=0, vmax=max_fail, cmap='Reds', integrate=True, window=window, arrayspace=USEARRSPACE) pplt.plot_xz(ax_slc(4), dump, dump['floors'] & 4, label="B_RHO", vmin=0, vmax=max_fail, cmap='Reds', integrate=True, window=window, arrayspace=USEARRSPACE) pplt.plot_xz(ax_slc(5), dump, dump['floors'] & 8, label="B_U", vmin=0, vmax=max_fail, cmap='Reds', integrate=True, window=window, arrayspace=USEARRSPACE) pplt.plot_xz(ax_slc(6), dump, dump['floors'] & 16, label="TEMP", vmin=0, vmax=max_fail, cmap='Reds', integrate=True, window=window, arrayspace=USEARRSPACE) pplt.plot_xz(ax_slc(7), dump, dump['floors'] & 32, label="GAMMA", vmin=0, vmax=max_fail, cmap='Reds', integrate=True, window=window, arrayspace=USEARRSPACE) pplt.plot_xz(ax_slc(8), dump, dump['floors'] & 64, label="KTOT", vmin=0, vmax=max_fail, cmap='Reds', integrate=True, window=window, arrayspace=USEARRSPACE) elif movie_type == "floors_old": ax_slc6 = lambda i: plt.subplot(2, 3, i) pplt.plot_slices(ax_slc6(1), ax_slc6(2), dump, 'log_rho', label=pretty('log_rho'), vmin=rho_l, vmax=rho_h, cmap='jet') max_fail = 1 pplt.plot_xz(ax_slc6(3), dump, dump['floors'] == 1, label="GEOM", vmin=0, vmax=max_fail, cmap='Reds') pplt.plot_xz(ax_slc6(4), dump, dump['floors'] == 2, label="SIGMA", vmin=0, vmax=max_fail, cmap='Reds') pplt.plot_xz(ax_slc6(5), dump, dump['floors'] == 3, label="GAMMA", vmin=0, vmax=max_fail, cmap='Reds') pplt.plot_xz(ax_slc6(6), dump, dump['floors'] == 4, label="KTOT", vmin=0, vmax=max_fail, cmap='Reds') else: # Strip global flags from the movie string l_movie_type = movie_type if "_ghost" in movie_type: l_movie_type = l_movie_type.replace("_ghost", "") if "_array" in l_movie_type: l_movie_type = l_movie_type.replace("_array", "") at = 0 if "_cross" in l_movie_type: l_movie_type = l_movie_type.replace("_cross", "") at = dump['n2'] // 2 if "_avg" in l_movie_type: l_movie_type = l_movie_type.replace("_avg", "") do_average = True else: do_average = False # Try to make a simple movie of just the stated variable # These are *informal*. Renormalize the colorscheme however we want #rho_l, rho_h = None, None if "_poloidal" in l_movie_type: ax = plt.subplot(1, 1, 1) var = l_movie_type.replace("_poloidal", "") pplt.plot_xz(ax, dump, var, at=at, label=pretty(var), vmin=rho_l, vmax=rho_h, window=window, arrayspace=USEARRSPACE, average=do_average, xlabel=False, ylabel=False, xticks=[], yticks=[], cbar=False, cmap='jet', field_overlay=False, shading=('gouraud', 'flat')[USEARRSPACE]) elif "_toroidal" in l_movie_type: ax = plt.subplot(1, 1, 1) var = l_movie_type.replace("_toroidal", "") pplt.plot_xy(ax, dump, var, at=at, label=pretty(var), vmin=rho_l, vmax=rho_h, window=window, arrayspace=USEARRSPACE, average=do_average, cbar=True, cmap='jet', shading=('gouraud', 'flat')[USEARRSPACE]) elif "_1d" in l_movie_type: ax = plt.subplot(1, 1, 1) var = l_movie_type.replace("_1d", "") ax.plot(dump['x'], dump[var][:, 0, 0], label=pretty(var)) ax.set_ylim((rho_l, rho_h)) ax.set_title(pretty(var)) else: ax_slc = [plt.subplot(1, 2, 1), plt.subplot(1, 2, 2)] ax = ax_slc[0] var = l_movie_type pplt.plot_slices(ax_slc[0], ax_slc[1], dump, var, at=at, label=pretty(l_movie_type), vmin=rho_l, vmax=rho_h, window=window, arrayspace=USEARRSPACE, average=do_average, cbar=True, cmap='jet', field_overlay=False, shading=('gouraud', 'flat')[USEARRSPACE]) # Labels if "divB" in movie_type: plt.suptitle(r"Max $\nabla \cdot B$ = {}".format( np.max(np.abs(dump['divB'])))) if "jsq" in movie_type: plt.subplots_adjust(hspace=0, wspace=0, left=0, right=1, bottom=0, top=1) if not USEARRSPACE: pplt.overlay_contours(ax, dump, 'sigma', [1]) #plt.subplots_adjust(left=0.03, right=0.97) plt.savefig(os.path.join(frame_dir, 'frame_%08d.png' % n), dpi=FIGDPI) plt.close(fig) del dump
def avg_dump(n): out = {} t = io.get_dump_time(dumps[n]) # When we don't know times, fudge # TODO accept -1 as "Not available" flag in the HDF5 spec if t == 0 and n != 0: t = 10 * n # Record out['coord/t'] = t if t < tstart or t > tend: # Still return the time return out print("Loading {} / {}: t = {}".format((n + 1), len(dumps), int(t)), file=sys.stderr) # TODO Add only what we need here... dump = pyHARM.load_dump(dumps[n], params=params, calc_derived=True, add_jcon=True, add_fails=True, add_floors=True) # Should we compute the time-averaged quantities? do_tavgs = (tavg_start <= t <= tavg_end) # EHT Radial profiles: Average only over the disk portion (excluding first & last pi/3 ~ "poles") if calc_ravgs: for var in [ 'rho', 'Pg', 'u^r', 'u^th', 'u^3', 'b^r', 'b^th', 'b^3', 'b', 'betainv', 'Ptot' ]: out['rt/' + var] = shell_avg(dump, var, j_slice=(jmin, jmax)) out['rt/' + var + '_notdisk'] = shell_avg(dump, var, j_slice=(0, jmin)) + \ shell_avg(dump, var, j_slice=(jmax, dump.header['n2'])) if do_tavgs: out['r/' + var] = out['rt/' + var] out['r/' + var + '_notdisk'] = out['rt/' + var + '_notdisk'] if calc_thavgs: if do_tavgs: # THETA AVERAGES for var in ['betainv', 'sigma']: out['th/' + var + '_25'] = theta_av(dump, var, i_of(r1d, 25), 5, fold=False) if calc_basic: # FIELD STRENGTHS # The HARM B_unit is sqrt(4pi)*c*sqrt(rho), and this is standard for EHT comparisons out['t/Phi_b'] = 0.5 * shell_sum( dump, np.fabs(dump['B1']), at_zone=iEH) # FLUXES # Radial profiles of Mdot and Edot, and their particular values # EHT code-comparison normalization has all these values positive for var, flux in [['Edot', 'FE'], ['Mdot', 'FM'], ['Ldot', 'FL']]: out['rt/' + flux] = shell_sum(dump, flux) if do_tavgs: out['r/' + flux] = shell_sum(dump, flux) out['t/' + var] = shell_sum(dump, flux, at_zone=iF) # Mdot and Edot are defined inward/positive at EH out['t/Mdot'] *= -1 out['t/Edot'] *= -1 if calc_diagnostics: # Maxima (for gauging floors) for var in ['sigma', 'betainv', 'Theta', 'U']: out['t/' + var + '_max'] = np.max(dump[var]) # Minima for var in ['rho', 'U']: out['t/' + var + '_min'] = np.min(dump[var]) out['rt/total_floors'] = np.sum(dump['floors'] != 0, axis=(1, 2)) out['rt/total_fails'] = np.sum(dump['fails'] != 0, axis=(1, 2)) if calc_phi: out['rt/Phi_b_sph'] = 0.5 * shell_sum(dump, np.fabs(dump['B1'])) out['rt/Phi_b_mid'] = np.zeros_like(out['rt/Phi_b_sph']) for i in range(out['rt/Phi_b_mid'].shape[0]): out['rt/Phi_b_mid'][i] = midplane_sum(dump, -dump['B2'], r_slice=(0, i)) if calc_madcc: out['rt/thrho'] = (shell_sum( dump, dump['rho'] * np.abs(np.pi / 2 - dump.grid.coords.th(dump.grid.coord_all()))) / shell_sum(dump, dump['rho'])) if do_tavgs: for var in [ 'rho', 'u^r', 'u^th', 'u^3', 'b^r', 'b^th', 'b^3', 'b', 'Pg', 'betainv', 'sigma' ]: out['rth/' + var] = dump[var].mean(axis=-1) if calc_madcc_optional: # Wavelength of fastest MRI mode for calculating suppression factor out['rt/lam_MRI'] = (shell_sum(dump, dump['rho'] * dump['lam_MRI']) / shell_sum(dump, dump['rho'])) # Correlation functions at specific radii for var in ['rho', 'betainv']: out['phit/' + var + '_cf10'] = corr_midplane(dump[var], at_i1=i_of(r1d, 10)) out['phit/' + var + '_cf20'] = corr_midplane(dump[var], at_i1=i_of(r1d, 20)) out['phit/' + var + '_cf30'] = corr_midplane(dump[var], at_i1=i_of(r1d, 30)) out['phit/' + var + '_cf50'] = corr_midplane(dump[var], at_i1=i_of(r1d, 50)) # Jet profile moments/ellipse for w_r in [50, 100]: for w_pole, w_slice in [('north', (0, jmin)), ('south', (jmax, dump.header['n2']))]: # CM out['thphit/jet_psi_' + w_pole + '_' + str(w_r)] = dump['jet_psi'][i_of(r1d, w_r), :, :] out['t/M_' + w_pole + '_' + str(w_r)] = M = shell_sum( dump, dump['jet_psi'] * np.cos(dump['th']), j_slice=w_slice, at_r=w_r) out['t/X_' + w_pole + '_' + str(w_r)] = X = 1 / M * shell_sum( dump, dump['x'] * dump['jet_psi'] * np.cos(dump['th']), j_slice=w_slice, at_r=w_r) out['t/Y_' + w_pole + '_' + str(w_r)] = Y = 1 / M * shell_sum( dump, dump['y'] * dump['jet_psi'] * np.cos(dump['th']), j_slice=w_slice, at_r=w_r) # Moments out['t/Ixx_' + w_pole + '_' + str(w_r)] = shell_sum( dump, (dump['x'] - X)**2 * dump['jet_psi'] * np.cos(dump['th']), j_slice=w_slice, at_r=w_r) out['t/Iyy_' + w_pole + '_' + str(w_r)] = shell_sum( dump, (dump['y'] - Y)**2 * dump['jet_psi'] * np.cos(dump['th']), j_slice=w_slice, at_r=w_r) out['t/Ixy_' + w_pole + '_' + str(w_r)] = shell_sum( dump, (dump['x'] - X) * (dump['y'] - Y) * dump['jet_psi'] * np.cos(dump['th']), j_slice=w_slice, at_r=w_r) del M, X, Y if do_tavgs: # Full midplane correlation function, time-averaged for var in ['rho', 'betainv']: out['rphi/' + var + '_cf'] = corr_midplane(dump[var]) # Polar profiles of different fluxes and variables if calc_jet_profile: for var in [ 'rho', 'bsq', 'b^r', 'b^th', 'b^3', 'u^r', 'u^th', 'u^3', 'FM', 'FE', 'FE_EM', 'FE_Fl', 'FL', 'FL_EM', 'FL_Fl', 'betagamma', 'Be_nob', 'Be_b' ]: out['tht/' + var + '_100'] = np.sum(dump[var][iBZ], axis=-1) if do_tavgs: out['th/' + var + '_100'] = out['tht/' + var + '_100'] out['thphi/' + var + '_100'] = dump[var][iBZ] #out['rth/' + var] = dump[var].mean(axis=-1) # Blandford-Znajek Luminosity L_BZ # This is a lot of luminosities! if calc_jet_cuts: # TODO cut on phi/t averages? -- needs 2-pass cut... cuts = { 'sigma1': lambda dump: (dump['sigma'] > 1), 'Be_b0': lambda dump: (dump['Be_b'] > 0.02), 'Be_b1': lambda dump: (dump['Be_b'] > 1), 'Be_nob0': lambda dump: (dump['Be_nob'] > 0.02), 'Be_nob1': lambda dump: (dump['Be_nob'] > 1), # 'mu1' : lambda dump : (dump['mu'] > 1), 'bg1': lambda dump: (dump['betagamma'] > 1.0), 'bg05': lambda dump: (dump['betagamma'] > 0.5), 'allp': lambda dump: (dump['FE'] > 0) } # Terminology: # LBZ = E&M energy only, any cut # Lj = full E flux, any cut # Ltot = Lj_allp = full luminosity wherever it is positive for lum, flux in [['LBZ', 'FE_EM'], ['Lj', 'FE']]: for cut in cuts.keys(): out['rt/' + lum + '_' + cut] = shell_sum(dump, flux, mask=cuts[cut](dump)) out['t/' + lum + '_' + cut] = out['rt/' + lum + '_' + cut][iBZ] if do_tavgs: out['r/' + lum + '_' + cut] = out['rt/' + lum + '_' + cut] else: # Use the default cut from Paper V # These are the powers for the MADCC is_jet = dump['Be_b'] > 1 for lum, flux in [['Mdot_jet', 'FM'], ['P_jet', 'FE'], ['P_EM_jet', 'FE_EM'], ['P_PAKE_jet', 'FE_PAKE'], ['P_EN_jet', 'FE_EN'], ['Area_jet', '1']]: out['rt/' + lum] = shell_sum(dump, flux, mask=is_jet) for lum, flux in [['Area_mag', '1']]: out['rt/' + lum] = shell_sum(dump, flux, mask=(dump['sigma'] > 1)) for var in [ 'rho', 'Pg', 'u^r', 'u^th', 'u^3', 'b^r', 'b^th', 'b^3', 'b', 'betainv', 'Ptot' ]: out['rt/' + var + '_jet'] = shell_avg(dump, var, mask=is_jet) del is_jet if calc_lumproxy: rho, Pg, B = dump['rho'], dump['Pg'], dump['b'] # See EHT code comparison paper j = rho**3 / Pg**2 * np.exp(-0.2 * (rho**2 / (B * Pg**2))**(1. / 3.)) out['rt/Lum'] = shell_sum(dump, j, j_slice=(jmin, jmax)) if calc_gridtotals: # Total energy and current, summed by shells to allow cuts on radius for tot_name, var_name in [['Etot', 'JE0'], ['Jsq_inv', 'jsq'], ['Jsq_loc', 'current']]: out['rt/' + tot_name] = shell_sum(dump, var_name) if calc_efluxes: # Conserved (maybe; in steady state) 2D energy flux for var in ['JE0', 'JE1', 'JE2']: out['rt/' + var] = shell_sum(dump, var) if do_tavgs: out['rth/' + var] = dump[var].mean(axis=-1) # Total outflowing portions of variables if calc_outfluxes: for name, var in [['outflow', 'FM'], ['outEflow', 'FE']]: var_tmp = dump[var] out['rt/' + name] = shell_sum(dump, var_tmp, mask=(var_tmp > 0)) if do_tavgs: out['r/' + name] = out['rt/' + name] if calc_pdfs: for var, pdf_range in [['betainv', [-3.5, 3.5]], ['rho', [-7, 1]]]: # TODO handle negatives, pass on the range & bins var_tmp = np.log10(dump[var]) out['pdft/' + var], _ = np.histogram( var_tmp, bins=pdf_nbins, range=pdf_range, weights=np.repeat(dump['gdet'], var_tmp.shape[2]).reshape(var_tmp.shape), density=True) del var_tmp if calc_omega_bz and do_tavgs: Fcov01, Fcov13 = Fcov(dump, 0, 1), Fcov(dump, 1, 3) Fcov02, Fcov23 = Fcov(dump, 0, 2), Fcov(dump, 2, 3) vr, vth, vphi = dump['u^1'] / dump['u^0'], dump['u^2'] / dump[ 'u^0'], dump['u^3'] / dump['u^0'] out['rhth/omega'] = np.zeros((hdr['n1'], hdr['n2'] // 2)) out['rhth/omega_alt_num'] = np.zeros((hdr['n1'], hdr['n2'] // 2)) out['rhth/omega_alt_den'] = np.zeros((hdr['n1'], hdr['n2'] // 2)) out['rhth/omega_alt'] = np.zeros((hdr['n1'], hdr['n2'] // 2)) out['rhth/vphi'] = np.zeros((hdr['n1'], hdr['n2'] // 2)) out['rhth/F13'] = np.zeros((hdr['n1'], hdr['n2'] // 2)) out['rhth/F01'] = np.zeros((hdr['n1'], hdr['n2'] // 2)) out['rhth/F23'] = np.zeros((hdr['n1'], hdr['n2'] // 2)) out['rhth/F02'] = np.zeros((hdr['n1'], hdr['n2'] // 2)) coord_hth = dump.grid.coord_all()[:, :, :hdr['n2'] // 2, 0] alpha_over_omega = dump.grid.lapse[Loci.CENT.value, :, :hdr['n2'] // 2] / (hdr['r_eh'] * np.sin( dump.grid.coords.th(coord_hth))) for i in range(hdr['n1']): out['rhth/F01'][i] = theta_av(dump, Fcov01, i, 1) out['rhth/F13'][i] = theta_av(dump, Fcov13, i, 1) out['rhth/F02'][i] = theta_av(dump, Fcov02, i, 1) out['rhth/F23'][i] = theta_av(dump, Fcov23, i, 1) out['rhth/omega'][i] = out['rhth/F01'][i] / out['rhth/F13'][i] out['rhth/omega_alt_num'][i] = theta_av( dump, vr * dump['B3'] * dump['B2'] + vth * dump['B3'] * dump['B1'], i, 1) out['rhth/omega_alt_den'][i] = theta_av(dump, dump['B2'] * dump['B1'], i, 1) out['rhth/omega_alt'][i] = theta_av( dump, vr * dump['B3'] / dump['B1'] + vth * dump['B3'] / dump['B2'], i, 1) out['rhth/vphi'][i] = theta_av(dump, vphi, i, 1) out['rhth/omega_alt'] *= -alpha_over_omega del Fcov01, Fcov13, vr, vth, vphi this_process = psutil.Process(os.getpid()) print("Memory use: {} GB".format(this_process.memory_info().rss / 10**9)) del dump return out
hdr = dump.header if dump['r'].ndim == 3: r1d = dump['r'][:, hdr['n2'] // 2, 0] elif dump['r'].ndim == 2: r1d = dump['r'][:, hdr['n2'] // 2] elif dump['r'].ndim == 1: r1d = dump['r'] jmin, jmax = get_eht_disk_j_vals(dump) del dump # Leave several extra zones if using MKS3 coordinates if hdr['coordinates'] == "mks3": iEH = i_of(r1d, hdr['r_eh']) + 4 else: iEH = i_of(r1d, hdr['r_eh']) # Measure fluxes at event horizon. TODO options here? iF = iEH # Max radius when computing "total" energy iEmax = i_of(r1d, 40) # BZ luminosity # 100M seems like the standard measuring spot (or at least, BHAC does it that way) # L_BZ seems constant* after that, but much higher within ~50M if hdr['r_out'] < 100 or r1d[ -1] < 100: # If in theory or practice the sim is small... if r1d[-1] > 40: