analistica/ex-1/plots/misc.py

#!/usr/bin/env python

from matplotlib import pyplot as plt
import numpy as np
import argparse
import sys


# -i  → agrs.show   → open plot in interactive mode
# -s  → args.save   → save plot
#
# -n  → agrs.note   → do the Landau plot with the notes
# -b  → args.both   → do the Landau-Moyal plot
# -m  → args.area   → do the Landau median plot
# -gm → args.gamo   → do the Gaussian-Moyal compare plot
# -gl → args.gala   → do the Gaussian-Landau compare plot
# -c  → args.cauchy → do the Cauchy PDF plot


# Moyal PDF
def moyal(x, μ, σ):
    N = (1)/(np.sqrt(2 * np.pi) * σ)
    return N * np.exp(- 0.5 * ((x - μ)/σ + np.exp(- (x - μ)/σ)))

# Cauchy PDF
def cauchy(x):
    return 1/(np.pi*(1 + x**2))

# Gauss PDF
def gauss(x, μ, σ):
    N = (1)/(np.sqrt(2 * np.pi) * σ)
    return N * np.exp(- 0.5 * ((x - μ)/σ)**2)


def main(args):
    # prepare figure
    if args.show:
        plt.figure()
    elif args.save:
        plt.rcParams['font.size'] = 8
        if (args.both or args.area or args.cauchy):
            plt.figure(figsize=(3, 2))
        elif (args.gamo or args.gala):
            plt.figure(figsize=(2.5, 2))
        elif args.note:
            plt.figure(figsize=(5, 3))

    if args.note:
        # useful coordinates
        y_min = -0.0086     # y min axes
        y_max = 0.1895      # y max axes
        me = -0.22          # mode
        f_me = 0.1806       # f(mode)
        h_f_me = f_me/2     # falf f(mode)
        x_m = -1.5867       # x₋
        x_p = 2.4330        # x₊

        # draw the lines
        plt.plot([-10, me],  [f_me, f_me],     color='gray')
        plt.plot([me, me],  [f_me, y_min],     color='gray')
        plt.plot([-10, x_p], [h_f_me, h_f_me], color='gray')
        plt.plot([x_m, x_m], [y_min, h_f_me],  color='gray')
        plt.plot([x_p, x_p], [y_min, h_f_me],  color='gray')

        # draw the notes
        s = 0.012
        S = 0.2
        plt.annotate('$f(m_e)$',   [-10 + S, f_me - s])
        plt.annotate('$f(m_e)/2$', [-10 + S, h_f_me - s])
        plt.annotate('$x_-$',      [x_m + S, y_min + s/2])
        plt.annotate('$x_+$',      [x_p + S, y_min + s/2])
        plt.annotate('$m_e$',      [me + S, y_min + s/2])

        # prepare plot
        plt.title('Landau distribution', loc='right')
        plt.xlim(-10, 10)
        plt.ylim(y_min, y_max)

    if (args.both):
        plt.ylim(-0.015, 0.265)

    # draw Landau plot
    if not args.gamo:
        x, y = np.loadtxt(sys.stdin, unpack=True)
        plt.plot(x, y, color='#92182b', label='Landau')

    # Color the area under the function
    if args.area:
        median = 1.3557804
        x0 = np.argmin(abs(median - x))
        plt.fill_between(x[:x0], y[:x0], color='#ff99a8')
        plt.fill_between(x[x0:], y[x0:], color='#b28e94')

    # do Cauchy plot
    if args.cauchy:
        x = np.arange(-10, 10, 0.01)
        plt.plot(x, cauchy(x), color='gray', label='Cauchy')
        plt.legend()

    μ = -0.22278298
    σ = 1.1191486

    # do Moyal plot
    if (args.both or args.gamo):
        x = np.arange(-10, 10, 0.01)
        plt.plot(x, moyal(x, μ, σ), color='gray', label='Moyal')
        plt.legend()

    # do Gaussian plot
    if (args.gamo or args.gala):
        x = np.arange(-10, 10, 0.01)
        if args.gala:
            a = 1.97
        if args.gamo:
            a = 1.65
        plt.plot(x, gauss(x, μ, a*σ), color='#cfb017', label='Gauss')
        plt.legend()

    # save figure
    plt.tight_layout()
    if args.show:
        plt.show()
    elif args.save:
        if args.note:
            plt.savefig('notes/images/1-notes.pdf')
        if args.both:
            plt.savefig('slides/images/both-pdf.pdf', transparent=True)
        if args.area:
            plt.savefig('slides/images/median.pdf', transparent=True)
        if args.cauchy:
            plt.savefig('slides/images/cauchy-pdf.pdf', transparent=True)
        if args.gamo:
            plt.savefig('slides/images/gamo-pdf.pdf')
        if args.gala:
            plt.savefig('slides/images/gala-pdf.pdf')


if __name__ == '__main__':
    # Parse data
    parser = argparse.ArgumentParser()
    parser.add_argument('-i', '--show',
                        action='store_true', default=False,
                        help='show the plots')
    parser.add_argument('-s', '--save',
                        action='store_true', default=False,
                        help='save the plots')
    parser.add_argument('-b', '--both',
                        action='store_true', default=False,
                        help='do the both plot')
    parser.add_argument('-n', '--note',
                        action='store_true', default=False,
                        help='do the note plot')
    parser.add_argument('-a', '--area',
                        action='store_true', default=False,
                        help='do the area plot')
    parser.add_argument('-c', '--cauchy',
                        action='store_true', default=False,
                        help='do the cauchy plot')
    parser.add_argument('-gl', '--gala',
                        action='store_true', default=False,
                        help='do the gaussian-landau plot')
    parser.add_argument('-gm', '--gamo',
                        action='store_true', default=False,
                        help='do the gaussian-moyal plot')
    args = parser.parse_args()
    main(args)