snippet_util.plot3d_energyspectrum

Sample to plot 3D for DifferentialFluxGrid.

"""Sample to plot 3D for :class:`DifferentialFluxGrid`.
"""

import numpy as np
from irfpy.util import energyspectrum

def plot3d(dflx):

    import matplotlib.pyplot as plt
    from mpl_toolkits.mplot3d import Axes3D

    fig = plt.figure()

    ax = fig.add_subplot(111, projection='3d')

    elist = np.log10(dflx.energy_grid)
    emin = int(elist.min()) - 1
    emax = int(elist.max()) + 1
    erng = emax - emin
    elist = elist - emin

    tlist = dflx.theta_grid
    plist = dflx.phi_grid

    e, t, p = np.meshgrid(elist, tlist, plist, indexing='ij')

    ex = e * np.sin(t) * np.cos(p)    # (96, 32, 64) shape
    ey = e * np.sin(t) * np.sin(p)
    ez = e * np.cos(t)

    jlist = np.log10(dflx.dflux)
    jmin = jlist.min()
    jmax = jlist.max()

    jmax = int(jmax) + 1
    jmin = jmax - 6

    clist = jlist.clip(jmin, jmax)
    clist = (clist - jmin) / (jmax - jmin) * 255.

#    slist = np.zeros_like(clist) + 20
#    slist = np.where(clist <= jmin, 0, 20)
    slist = clist / 255. * 50.

#    print(jmin, jmax, clist.min(), clist.max())

    ax.scatter(ex, ey, ez, c=clist, s=slist, alpha=0.1, edgecolor='none')
    ax.set_xlabel("log(Ex) - {} [eV]".format(emin))
    ax.set_ylabel("log(Ey) - {} [eV]".format(emin))
    ax.set_zlabel("log(Ez) - {} [eV]".format(emin))

#    fig.savefig('plot3d_energyspectrum.png')


def main():
    elist = np.logspace(1, 4, 32)   # 10eV to 10000eV.
    tlist = np.linspace(0, np.pi, 17)
    tlist = (tlist[1:] + tlist[:-1]) / 2.    # 16 steps for 180 deg.
    plist = np.linspace(0, 2 * np.pi, 33)[:-1]  # 32 steps for 360 deg (The last one should not be taken into acount to avoid duplication (2pi = 0)

    dflx = energyspectrum.maxwellDifferentialFluxGrid(1e6, [400e3, 0, 0], 80e3,
                elist, tlist, plist)

#    plot3d(dflx)
    ax, sct = dflx.plot3d(alpha=0.1)

    return ax, sct

if __name__ == '__main__':
    retval = main()