appl121105_iotorus.plotter

A general routine of plotting.

For debuging, simple to use plotting routine. No detailed setting is possible, but just call contour_xz().

''' A general routine of plotting.

For debuging, simple to use plotting routine.
No detailed setting is possible, but just call
:func:`contour_xz`.

'''
import os
import sys
import logging
logging.basicConfig()
import datetime
import math

import matplotlib.pyplot as plt
import numpy as np
import scipy as sp

def contour_xz(data_callable, xgrids, zgrids, yval=0., fill=True, fig=None, ax=None, scale='linear', *args, **kwds):
    '''

    :param data_callable: A callable of the data
    :param xgrids: Grid in x coordinates
    :param zgrids: Grid in z coordinates
    :keyword fig: Figure instance. If ``None``, newly created.
    :keyword ax: Axis instance. If ``None``, newly created by subplot(111).
    :returns: (fig, ax)
    '''
    if fig == None:
        fig = plt.figure()
    if ax == None:
        ax = fig.add_subplot(111)

    xx, zz = np.meshgrid(xgrids, zgrids)

    fld = np.zeros_like(xx)

    for ix in range(len(xgrids)):
        for iz in range(len(zgrids)):
            fld[iz, ix] = data_callable(xx[iz, ix], yval, zz[iz, ix])

    if scale == 'log':
        fld = np.log10(fld)
    
    if fill:
        f = ax.contourf(xx, zz, fld, *args, **kwds)
        fig.colorbar(f)

    else:
        ax.contour(xx, zz, fld, *args, **kwds)

    ax.set_xlabel('X')
    ax.set_ylabel('Z')

    return (fig, ax)

if __name__ == "__main__":
    from pyana.pep import iotorus
#    n = iotorus.TwoPeakModel()
#    n = iotorus.PlasmaTorusEnvdocL()
#    n = iotorus.PlasmaTorusEnvdocLSech6(idx=10)
    inputscale = 1.
    colorscale = 0.

    from . import plasmamodel
#    n = plasmamodel.SimpleCorotationalFlow2()
    n = plasmamodel.CorotationalFlowMJconv()
    inputscale = 71492.
    colorscale = 6.

    densityfunc = lambda x, y, z: n.get_density(x, y, z, sun_long_sys3=[90.])
    
    fig, ax = contour_xz(densityfunc, np.arange(-10, 10, 0.5) * inputscale, np.arange(-10, 10, 0.5) * inputscale, scale='log', levels=np.array([0.5, 1, 1.5, 2, 2.5, 3, 3.5, 4, 4.5, 5]) + colorscale)
    fig.savefig('plotter.png')
    
    fig, ax = contour_xz(n.get_density, np.arange(5, 10, 0.1) * inputscale, np.arange(-2, 2, 0.1) * inputscale, scale='log', levels=np.array([1.4, 1.7, 2, 2.3, 2.6, 2.9, 3.2, 3.5]) + colorscale)
    fig.savefig('plotter5_10.png')