
import matplotlib.pylab as plt
from mpl_toolkits.axes_grid.inset_locator import inset_axes
import numpy as np
from . import probscores

def plot_reldiag(reldiag, ax=None):
    """Plot a reliability diagram.
    
    Parameters
    ----------
    reldiag : dict
        A ROC curve object created by probscores.reldiag_init.
    ax : axis handle
        Axis handle for the figure. If set to None, the handle is taken from 
        the current figure (matplotlib.pylab.gca()).
    """
    if ax is None:
        ax = plt.gca()
    
    # Plot the reliability diagram.
    # TODO: Use the compute method here.
    f = 1.0 * reldiag["Y_sum"] / reldiag["num_idx"]
    r = 1.0 * reldiag["X_sum"] / reldiag["num_idx"]
    
    mask = np.logical_and(np.isfinite(r), np.isfinite(f))
    
    ax.plot(r[mask], f[mask], "kD-")
    ax.plot([0, 1], [0, 1], "k--")
    
    ax.set_xlim(0, 1)
    ax.set_ylim(0, 1)
    
    ax.grid(True, ls=':')
    
    ax.set_xlabel("Forecast probability")
    ax.set_ylabel("Observed relative frequency")
    
    # Plot sharpness diagram into an inset figure.
    iax = inset_axes(ax, width="35%", height="20%", loc=4, borderpad=3.5)
    bw = reldiag["bin_edges"][2] - reldiag["bin_edges"][1]
    iax.bar(reldiag["bin_edges"][:-1], reldiag["sample_size"], width=bw, 
            align="edge", color="gray", edgecolor="black")
    iax.set_yscale("log", basey=10)
    iax.set_xticks(reldiag["bin_edges"])
    iax.set_xticklabels(["%.1f" % max(v, 1e-6) for v in reldiag["bin_edges"]])
    yt_min = int(max(np.floor(np.log10(min(reldiag["sample_size"][:-1]))), 1))
    yt_max = int(np.ceil(np.log10(max(reldiag["sample_size"][:-1]))))
    t = [pow(10.0, k) for k in range(yt_min, yt_max)]
    
    iax.set_yticks([int(t_) for t_ in t])
    iax.set_xlim(0.0, 1.0)
    iax.set_ylim(t[0], 5*t[-1])
    iax.set_ylabel("log10(samples)")
    iax.yaxis.tick_right()
    iax.yaxis.set_label_position("right")
    iax.tick_params(axis="both", which="major", labelsize=6)

def plot_ROC(ROC, ax=None):
    """Plot a ROC curve.
    
    Parameters
    ----------
    ROC : dict
        A ROC curve object created by probscores.ROC_curve_init.
    ax : axis handle
        Axis handle for the figure. If set to None, the handle is taken from 
        the current figure (matplotlib.pylab.gca()).
    """
    if ax is None:
        ax = plt.gca()
    
    POFD,POD,area = probscores.ROC_curve_compute(ROC, compute_area=True)
    p_thr = ROC["prob_thrs"]
    
    ax.plot([0, 1], [0, 1], "k--")
    ax.set_xlim(0, 1)
    ax.set_ylim(0, 1)
    ax.set_xlabel("False alarm rate (POFD)")
    ax.set_ylabel("Probability of detection (POD)")
    ax.grid(True, ls=':')
    
    l, = ax.plot(POFD, POD, "kD-")
    opt_prob_thr_idx = np.argmax(np.array(POD) - np.array(POFD))
    ax.scatter([POFD[opt_prob_thr_idx]], [POD[opt_prob_thr_idx]], c='r', s=150, 
               facecolors=None, edgecolors='r')
    
    for p_thr_,x,y in zip(p_thr, POFD, POD):
        if p_thr_ > 0.05 and p_thr_ < 0.95:
            ax.text(x+0.02, y-0.02, "%.2f" % p_thr_, fontsize=7)


@@ -118,26 +118,939 @@
  import 
-probscores
+ensscores, probscores%0A%0Adef plot_rankhist(rankhist, ax=None):%0A    %22%22%22Plot a rank histogram.%0A    %0A    Parameters%0A    ----------%0A    rankhist : dict%0A        A rank histogram object created by ensscores.rankhist_init.%0A    ax : axis handle%0A        Axis handle for the figure. If set to None, the handle is taken from %0A        the current figure (matplotlib.pylab.gca()).%0A    %22%22%22%0A    if ax is None:%0A        ax = plt.gca()%0A    %0A    r = ensscores.rankhist_compute(rankhist)%0A    x = np.linspace(0, 1, rankhist%5B%22num_ens_members%22%5D + 1)%0A    ax.bar(x, r, width=1.0/len(x), align=%22edge%22, color=%22gray%22, edgecolor=%22black%22)%0A    %0A    ax.set_xticks(x%5B::3%5D + (x%5B1%5D - x%5B0%5D))%0A    ax.set_xticklabels(np.arange(1, len(x))%5B::3%5D)%0A    ax.set_xlim(0, 1+1.0/len(x))%0A    ax.set_ylim(0, np.max(r)*1.25)%0A    %0A    ax.set_xlabel(%22Rank of observation (among ensemble members)%22)%0A    ax.set_ylabel(%22Relative frequency%22)%0A    %0A    ax.grid(True, axis='y', ls=':')
 %0A%0Adef pl
@@ -1434,24 +1434,24 @@
 .gca()%0A    %0A
+
     # Plot t
@@ -1478,49 +1478,8 @@
 am.%0A
-    # TODO: Use the compute method here.%0A
     

Add plotting of rank histograms.