Exemplo n.º 1
0
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'])
Exemplo n.º 2
0
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))
Exemplo n.º 3
0
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)
Exemplo n.º 4
0
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
Exemplo n.º 5
0
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
Exemplo n.º 6
0
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
Exemplo n.º 7
0
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
Exemplo n.º 8
0
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: