#! /usr/bin/env python3
"""
    GPPDF plotter

    This script takes the output of `gppdf.py` (a numpy archive) and plots the
    resulting samples, flavour by PDF flavour. The output is written to file at
    'plt_{setname}.pdf', where `setname` is the GPPDF set name.

"""
from matplotlib.backends.backend_pdf import PdfPages
import matplotlib.pyplot as plt
from matplotlib import gridspec
from matplotlib.ticker import MaxNLocator

from sys import argv
import numpy as np
import datetime
import argparse

# Matplotlib style
plt.style.use('seaborn-colorblind')

# Number of active flavours at initial scale
labels = {-6: "tbar", -5: "bbar", -4: "cbar", -3: "sbar", -2: "ubar", -1: "dbar",
          21: "g", 1: "d", 2: "u", 3: "s", 4: "c", 5: "b", 6: "t"}

def setup_figure(xmin):
    """ Setup a split log-lin figure over [xmin, 1] for plotting PDFs.
        Returns (figure, log-axis, lin-axis) """
    # Setup gridspec
    gs = gridspec.GridSpec(1, 2)
    gs.update(wspace=0.00, hspace=0.00)

    # Setup figure
    w, h = plt.figaspect(0.5)
    fig   = plt.figure(figsize=(w, h))
    logax = fig.add_subplot(gs[0])
    linax = fig.add_subplot(gs[1])

    # Axis formatting
    linax.set_xlim([0.1, 1])
    logax.set_xlim([min(gpdata["xgrid"]), 0.1])
    plt.setp(linax.get_yticklabels(), visible=False)
    logax.set_xscale('log')
    linax.xaxis.set_major_locator(MaxNLocator(5, prune='lower'))
    return fig, logax, linax

def gppdf_figure(fig, gpdata, ipdf):
    """ Given a split figure `fig` generated by `setup_figure`, and a
    Gaussian-process sample `gpdata`, plots the `ipdf`th PDF flavour to `fig`."""

    xgrid = gpdata["xgrid"]
    nxpts = len(xgrid)
    nsamples = len(gpdata["samples"])

    # GP Uncertainty
    error = np.sqrt(np.diagonal(gpdata["covariance"]))

    # GP Mean and standard-deviation
    start_index = nxpts*ipdf
    end_index   = start_index + nxpts
    mslice = gpdata["mean"][start_index:end_index]
    eslice = error[start_index:end_index]

    for ax in fig.get_axes():
        ax.grid()
        for irep in range(0, nsamples):
            gpslice = gpdata["samples"][irep][start_index:end_index]
            if irep == 0:
                ax.plot(xgrid, gpslice, alpha=0.05, color='b', label="GP Samples")
            else:
                ax.plot(xgrid, gpslice, alpha=0.05, color='b')
        ax.plot(xgrid, mslice+eslice, color='black', linestyle='--')
        ax.plot(xgrid, mslice-eslice, color='black', linestyle='--', label="GP 1-sigma")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("gppdf", help="GPPDF Sample (npz archive)", type=str)
    args = parser.parse_args()
    target = args.gppdf

    print(f"Loading archive {target}")
    gpdata = np.load(target)

    # Plot PDFs
    with PdfPages(f'plt_{gpdata["setname"]}.pdf') as output:
        d = output.infodict()
        d['Title'] = 'GPPDF plots'
        d['Author'] = 'Nathan Hartland'
        d['CreationDate'] = datetime.datetime.today()
        # Make one figure per PDF flavour
        for ipdf, pdf in enumerate(gpdata["flavours"]):
            xmin = min(gpdata["xgrid"])
            fig, logax, linax = setup_figure(xmin)
            gppdf_figure(fig, gpdata, ipdf)

            # Add legend to linear-axis
            leg = linax.legend(loc='best')

            # Labels and title
            logax.set_ylabel(f'$x${labels[pdf]}(x, Q={gpdata["Q0"]})')
            fig.suptitle(f'Gaussian process sample of {labels[pdf]}(x, Q={gpdata["Q0"]}) from {gpdata["prior"]}')

            output.savefig(fig)

6ef51856ae2cdc2d0adddf550111ba606b1fa215
<mask token>


def setup_figure(xmin):
    """ Setup a split log-lin figure over [xmin, 1] for plotting PDFs.
        Returns (figure, log-axis, lin-axis) """
    gs = gridspec.GridSpec(1, 2)
    gs.update(wspace=0.0, hspace=0.0)
    w, h = plt.figaspect(0.5)
    fig = plt.figure(figsize=(w, h))
    logax = fig.add_subplot(gs[0])
    linax = fig.add_subplot(gs[1])
    linax.set_xlim([0.1, 1])
    logax.set_xlim([min(gpdata['xgrid']), 0.1])
    plt.setp(linax.get_yticklabels(), visible=False)
    logax.set_xscale('log')
    linax.xaxis.set_major_locator(MaxNLocator(5, prune='lower'))
    return fig, logax, linax


def gppdf_figure(fig, gpdata, ipdf):
    """ Given a split figure `fig` generated by `setup_figure`, and a
    Gaussian-process sample `gpdata`, plots the `ipdf`th PDF flavour to `fig`."""
    xgrid = gpdata['xgrid']
    nxpts = len(xgrid)
    nsamples = len(gpdata['samples'])
    error = np.sqrt(np.diagonal(gpdata['covariance']))
    start_index = nxpts * ipdf
    end_index = start_index + nxpts
    mslice = gpdata['mean'][start_index:end_index]
    eslice = error[start_index:end_index]
    for ax in fig.get_axes():
        ax.grid()
        for irep in range(0, nsamples):
            gpslice = gpdata['samples'][irep][start_index:end_index]
            if irep == 0:
                ax.plot(xgrid, gpslice, alpha=0.05, color='b', label=
                    'GP Samples')
            else:
                ax.plot(xgrid, gpslice, alpha=0.05, color='b')
        ax.plot(xgrid, mslice + eslice, color='black', linestyle='--')
        ax.plot(xgrid, mslice - eslice, color='black', linestyle='--',
            label='GP 1-sigma')


<mask token>

<mask token>
plt.style.use('seaborn-colorblind')
<mask token>


def setup_figure(xmin):
    """ Setup a split log-lin figure over [xmin, 1] for plotting PDFs.
        Returns (figure, log-axis, lin-axis) """
    gs = gridspec.GridSpec(1, 2)
    gs.update(wspace=0.0, hspace=0.0)
    w, h = plt.figaspect(0.5)
    fig = plt.figure(figsize=(w, h))
    logax = fig.add_subplot(gs[0])
    linax = fig.add_subplot(gs[1])
    linax.set_xlim([0.1, 1])
    logax.set_xlim([min(gpdata['xgrid']), 0.1])
    plt.setp(linax.get_yticklabels(), visible=False)
    logax.set_xscale('log')
    linax.xaxis.set_major_locator(MaxNLocator(5, prune='lower'))
    return fig, logax, linax


def gppdf_figure(fig, gpdata, ipdf):
    """ Given a split figure `fig` generated by `setup_figure`, and a
    Gaussian-process sample `gpdata`, plots the `ipdf`th PDF flavour to `fig`."""
    xgrid = gpdata['xgrid']
    nxpts = len(xgrid)
    nsamples = len(gpdata['samples'])
    error = np.sqrt(np.diagonal(gpdata['covariance']))
    start_index = nxpts * ipdf
    end_index = start_index + nxpts
    mslice = gpdata['mean'][start_index:end_index]
    eslice = error[start_index:end_index]
    for ax in fig.get_axes():
        ax.grid()
        for irep in range(0, nsamples):
            gpslice = gpdata['samples'][irep][start_index:end_index]
            if irep == 0:
                ax.plot(xgrid, gpslice, alpha=0.05, color='b', label=
                    'GP Samples')
            else:
                ax.plot(xgrid, gpslice, alpha=0.05, color='b')
        ax.plot(xgrid, mslice + eslice, color='black', linestyle='--')
        ax.plot(xgrid, mslice - eslice, color='black', linestyle='--',
            label='GP 1-sigma')


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('gppdf', help='GPPDF Sample (npz archive)', type=str)
    args = parser.parse_args()
    target = args.gppdf
    print(f'Loading archive {target}')
    gpdata = np.load(target)
    with PdfPages(f"plt_{gpdata['setname']}.pdf") as output:
        d = output.infodict()
        d['Title'] = 'GPPDF plots'
        d['Author'] = 'Nathan Hartland'
        d['CreationDate'] = datetime.datetime.today()
        for ipdf, pdf in enumerate(gpdata['flavours']):
            xmin = min(gpdata['xgrid'])
            fig, logax, linax = setup_figure(xmin)
            gppdf_figure(fig, gpdata, ipdf)
            leg = linax.legend(loc='best')
            logax.set_ylabel(f"$x${labels[pdf]}(x, Q={gpdata['Q0']})")
            fig.suptitle(
                f"Gaussian process sample of {labels[pdf]}(x, Q={gpdata['Q0']}) from {gpdata['prior']}"
                )
            output.savefig(fig)

<mask token>
plt.style.use('seaborn-colorblind')
labels = {(-6): 'tbar', (-5): 'bbar', (-4): 'cbar', (-3): 'sbar', (-2):
    'ubar', (-1): 'dbar', (21): 'g', (1): 'd', (2): 'u', (3): 's', (4): 'c',
    (5): 'b', (6): 't'}


def setup_figure(xmin):
    """ Setup a split log-lin figure over [xmin, 1] for plotting PDFs.
        Returns (figure, log-axis, lin-axis) """
    gs = gridspec.GridSpec(1, 2)
    gs.update(wspace=0.0, hspace=0.0)
    w, h = plt.figaspect(0.5)
    fig = plt.figure(figsize=(w, h))
    logax = fig.add_subplot(gs[0])
    linax = fig.add_subplot(gs[1])
    linax.set_xlim([0.1, 1])
    logax.set_xlim([min(gpdata['xgrid']), 0.1])
    plt.setp(linax.get_yticklabels(), visible=False)
    logax.set_xscale('log')
    linax.xaxis.set_major_locator(MaxNLocator(5, prune='lower'))
    return fig, logax, linax


def gppdf_figure(fig, gpdata, ipdf):
    """ Given a split figure `fig` generated by `setup_figure`, and a
    Gaussian-process sample `gpdata`, plots the `ipdf`th PDF flavour to `fig`."""
    xgrid = gpdata['xgrid']
    nxpts = len(xgrid)
    nsamples = len(gpdata['samples'])
    error = np.sqrt(np.diagonal(gpdata['covariance']))
    start_index = nxpts * ipdf
    end_index = start_index + nxpts
    mslice = gpdata['mean'][start_index:end_index]
    eslice = error[start_index:end_index]
    for ax in fig.get_axes():
        ax.grid()
        for irep in range(0, nsamples):
            gpslice = gpdata['samples'][irep][start_index:end_index]
            if irep == 0:
                ax.plot(xgrid, gpslice, alpha=0.05, color='b', label=
                    'GP Samples')
            else:
                ax.plot(xgrid, gpslice, alpha=0.05, color='b')
        ax.plot(xgrid, mslice + eslice, color='black', linestyle='--')
        ax.plot(xgrid, mslice - eslice, color='black', linestyle='--',
            label='GP 1-sigma')


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('gppdf', help='GPPDF Sample (npz archive)', type=str)
    args = parser.parse_args()
    target = args.gppdf
    print(f'Loading archive {target}')
    gpdata = np.load(target)
    with PdfPages(f"plt_{gpdata['setname']}.pdf") as output:
        d = output.infodict()
        d['Title'] = 'GPPDF plots'
        d['Author'] = 'Nathan Hartland'
        d['CreationDate'] = datetime.datetime.today()
        for ipdf, pdf in enumerate(gpdata['flavours']):
            xmin = min(gpdata['xgrid'])
            fig, logax, linax = setup_figure(xmin)
            gppdf_figure(fig, gpdata, ipdf)
            leg = linax.legend(loc='best')
            logax.set_ylabel(f"$x${labels[pdf]}(x, Q={gpdata['Q0']})")
            fig.suptitle(
                f"Gaussian process sample of {labels[pdf]}(x, Q={gpdata['Q0']}) from {gpdata['prior']}"
                )
            output.savefig(fig)

<mask token>
from matplotlib.backends.backend_pdf import PdfPages
import matplotlib.pyplot as plt
from matplotlib import gridspec
from matplotlib.ticker import MaxNLocator
from sys import argv
import numpy as np
import datetime
import argparse
plt.style.use('seaborn-colorblind')
labels = {(-6): 'tbar', (-5): 'bbar', (-4): 'cbar', (-3): 'sbar', (-2):
    'ubar', (-1): 'dbar', (21): 'g', (1): 'd', (2): 'u', (3): 's', (4): 'c',
    (5): 'b', (6): 't'}


def setup_figure(xmin):
    """ Setup a split log-lin figure over [xmin, 1] for plotting PDFs.
        Returns (figure, log-axis, lin-axis) """
    gs = gridspec.GridSpec(1, 2)
    gs.update(wspace=0.0, hspace=0.0)
    w, h = plt.figaspect(0.5)
    fig = plt.figure(figsize=(w, h))
    logax = fig.add_subplot(gs[0])
    linax = fig.add_subplot(gs[1])
    linax.set_xlim([0.1, 1])
    logax.set_xlim([min(gpdata['xgrid']), 0.1])
    plt.setp(linax.get_yticklabels(), visible=False)
    logax.set_xscale('log')
    linax.xaxis.set_major_locator(MaxNLocator(5, prune='lower'))
    return fig, logax, linax


def gppdf_figure(fig, gpdata, ipdf):
    """ Given a split figure `fig` generated by `setup_figure`, and a
    Gaussian-process sample `gpdata`, plots the `ipdf`th PDF flavour to `fig`."""
    xgrid = gpdata['xgrid']
    nxpts = len(xgrid)
    nsamples = len(gpdata['samples'])
    error = np.sqrt(np.diagonal(gpdata['covariance']))
    start_index = nxpts * ipdf
    end_index = start_index + nxpts
    mslice = gpdata['mean'][start_index:end_index]
    eslice = error[start_index:end_index]
    for ax in fig.get_axes():
        ax.grid()
        for irep in range(0, nsamples):
            gpslice = gpdata['samples'][irep][start_index:end_index]
            if irep == 0:
                ax.plot(xgrid, gpslice, alpha=0.05, color='b', label=
                    'GP Samples')
            else:
                ax.plot(xgrid, gpslice, alpha=0.05, color='b')
        ax.plot(xgrid, mslice + eslice, color='black', linestyle='--')
        ax.plot(xgrid, mslice - eslice, color='black', linestyle='--',
            label='GP 1-sigma')


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('gppdf', help='GPPDF Sample (npz archive)', type=str)
    args = parser.parse_args()
    target = args.gppdf
    print(f'Loading archive {target}')
    gpdata = np.load(target)
    with PdfPages(f"plt_{gpdata['setname']}.pdf") as output:
        d = output.infodict()
        d['Title'] = 'GPPDF plots'
        d['Author'] = 'Nathan Hartland'
        d['CreationDate'] = datetime.datetime.today()
        for ipdf, pdf in enumerate(gpdata['flavours']):
            xmin = min(gpdata['xgrid'])
            fig, logax, linax = setup_figure(xmin)
            gppdf_figure(fig, gpdata, ipdf)
            leg = linax.legend(loc='best')
            logax.set_ylabel(f"$x${labels[pdf]}(x, Q={gpdata['Q0']})")
            fig.suptitle(
                f"Gaussian process sample of {labels[pdf]}(x, Q={gpdata['Q0']}) from {gpdata['prior']}"
                )
            output.savefig(fig)

#! /usr/bin/env python3
"""
    GPPDF plotter

    This script takes the output of `gppdf.py` (a numpy archive) and plots the
    resulting samples, flavour by PDF flavour. The output is written to file at
    'plt_{setname}.pdf', where `setname` is the GPPDF set name.

"""
from matplotlib.backends.backend_pdf import PdfPages
import matplotlib.pyplot as plt
from matplotlib import gridspec
from matplotlib.ticker import MaxNLocator

from sys import argv
import numpy as np
import datetime
import argparse

# Matplotlib style
plt.style.use('seaborn-colorblind')

# Number of active flavours at initial scale
labels = {-6: "tbar", -5: "bbar", -4: "cbar", -3: "sbar", -2: "ubar", -1: "dbar",
          21: "g", 1: "d", 2: "u", 3: "s", 4: "c", 5: "b", 6: "t"}

def setup_figure(xmin):
    """ Setup a split log-lin figure over [xmin, 1] for plotting PDFs.
        Returns (figure, log-axis, lin-axis) """
    # Setup gridspec
    gs = gridspec.GridSpec(1, 2)
    gs.update(wspace=0.00, hspace=0.00)

    # Setup figure
    w, h = plt.figaspect(0.5)
    fig   = plt.figure(figsize=(w, h))
    logax = fig.add_subplot(gs[0])
    linax = fig.add_subplot(gs[1])

    # Axis formatting
    linax.set_xlim([0.1, 1])
    logax.set_xlim([min(gpdata["xgrid"]), 0.1])
    plt.setp(linax.get_yticklabels(), visible=False)
    logax.set_xscale('log')
    linax.xaxis.set_major_locator(MaxNLocator(5, prune='lower'))
    return fig, logax, linax

def gppdf_figure(fig, gpdata, ipdf):
    """ Given a split figure `fig` generated by `setup_figure`, and a
    Gaussian-process sample `gpdata`, plots the `ipdf`th PDF flavour to `fig`."""

    xgrid = gpdata["xgrid"]
    nxpts = len(xgrid)
    nsamples = len(gpdata["samples"])

    # GP Uncertainty
    error = np.sqrt(np.diagonal(gpdata["covariance"]))

    # GP Mean and standard-deviation
    start_index = nxpts*ipdf
    end_index   = start_index + nxpts
    mslice = gpdata["mean"][start_index:end_index]
    eslice = error[start_index:end_index]

    for ax in fig.get_axes():
        ax.grid()
        for irep in range(0, nsamples):
            gpslice = gpdata["samples"][irep][start_index:end_index]
            if irep == 0:
                ax.plot(xgrid, gpslice, alpha=0.05, color='b', label="GP Samples")
            else:
                ax.plot(xgrid, gpslice, alpha=0.05, color='b')
        ax.plot(xgrid, mslice+eslice, color='black', linestyle='--')
        ax.plot(xgrid, mslice-eslice, color='black', linestyle='--', label="GP 1-sigma")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("gppdf", help="GPPDF Sample (npz archive)", type=str)
    args = parser.parse_args()
    target = args.gppdf

    print(f"Loading archive {target}")
    gpdata = np.load(target)

    # Plot PDFs
    with PdfPages(f'plt_{gpdata["setname"]}.pdf') as output:
        d = output.infodict()
        d['Title'] = 'GPPDF plots'
        d['Author'] = 'Nathan Hartland'
        d['CreationDate'] = datetime.datetime.today()
        # Make one figure per PDF flavour
        for ipdf, pdf in enumerate(gpdata["flavours"]):
            xmin = min(gpdata["xgrid"])
            fig, logax, linax = setup_figure(xmin)
            gppdf_figure(fig, gpdata, ipdf)

            # Add legend to linear-axis
            leg = linax.legend(loc='best')

            # Labels and title
            logax.set_ylabel(f'$x${labels[pdf]}(x, Q={gpdata["Q0"]})')
            fig.suptitle(f'Gaussian process sample of {labels[pdf]}(x, Q={gpdata["Q0"]}) from {gpdata["prior"]}')

            output.savefig(fig)
