distribution_function.py

#!/usr/bin/env python3

import shapefile
import numpy as np
import matplotlib.pyplot as plt
import itertools
from os import path

def plot_distr_funct(x, label = None, ax = None):
    """Sort and plot distribution function. x should be a Numpy array."""
    
    x_sorted = np.sort(x)
    nx = np.size(x)
    
    if nx != 0:
        if ax is None: fig, ax = plt.subplots()
        ax.plot(x_sorted, (1 + np.arange(nx)) / nx, label = label)
        print("minimum value:", x_sorted[0])
        print("maximum value:", x_sorted[- 1])
        
    print("number of values:", nx, "\n")

def read(dirname):
    """Read the three dbf files in dirname and return speed, radius and
    amplitude of outermost contour, radius and amplitude of maximum
    speed contour, as Numpy arrays.

    Select valid speed values.

    """
    extr_file = path.join(dirname, "extremum")
    outer_file = path.join(dirname, "outermost_contour")
    max_speed_file = path.join(dirname, "max_speed_contour")
    
    with shapefile.Reader(extr_file) as extremum, \
         shapefile.Reader(outer_file) as outerm_cont, \
         shapefile.Reader(max_speed_file) as max_speed_cont:
        speed = []
        rad_outer = []
        rad_speed = []
        amp_outer = []
        amp_speed = []

        for rec_extr, rec_outer, rec_max in zip(extremum.iterRecords(),
                                                outerm_cont.iterRecords(),
                                                max_speed_cont.iterRecords()):
            if rec_extr.speed != 1e4: speed.append(rec_extr.speed)

            rad_outer.append(rec_outer.r_eq_area)
            amp_outer.append(rec_extr.ssh - rec_outer.ssh)
                
            if rec_max.r_eq_area != - 100:
                rad_speed.append(rec_max.r_eq_area)
                amp_speed.append(rec_extr.ssh - rec_max.ssh)

    return {"speed": np.array(speed), "rad_outer": np.array(rad_outer),
            "rad_speed": np.array(rad_speed),
            "amp_outer": np.abs(np.array(amp_outer)),
            "amp_speed": np.abs(np.array(amp_speed))}

def fig_hist(xlabel, x, label):
    """Creates a complete figure with histogram. x and label can also be
    sequences."""
    
    fig, ax = plt.subplots()
    ax.hist(x, bins = "auto", histtype = "step", label = label)
    ax.set_xlabel(xlabel)
    ax.set_ylabel("number of occurrences")
    if label is not None: ax.legend()

def fig_distr_funct(xlabel, x, label = None):
    """Creates a complete figure with distribution function. x must be a
    sequence of Numpy arrays and label None or a sequence of strings.

    """
    
    fig, ax = plt.subplots()
    print(xlabel + ":")

    label_None = label is None
    if label_None: label = itertools.repeat(None)
        
    for xi, li in zip(x, label):
        plot_distr_funct(xi, li, ax)
        
    ax.set_xlabel(xlabel)
    ax.set_ylabel("distribution function")
    if not label_None: ax.legend()
    
def plot_all(dict_list, label = None):
    """dict_list: list of dictionaries. label: list of labels, one label
    for each dictionary.

    """
    
    fig_hist("speed, in m s-1", [d["speed"] for d in dict_list], label)
    fig_distr_funct("speed, when positive, in m s-1",
                    [d["speed"][d["speed"] >= 0] for d in dict_list], label)
    fig_distr_funct("speed, when negative, in m s-1",
                    [- d["speed"][d["speed"] < 0] for d in dict_list], label)
    
    fig_hist("equivalent radius of outermost contour, in km",
             [d["rad_outer"] for d in dict_list], label)
    fig_distr_funct("equivalent radius of outermost contour, in km",
                    [d["rad_outer"] for d in dict_list], label)

    fig_hist("equivalent radius of max-speed contour, in km",
             [d["rad_speed"] for d in dict_list], label)
    fig_distr_funct("equivalent radius of max-speed contour, in km",
                    [d["rad_speed"] for d in dict_list], label)

    fig_hist("amplitude of outermost contour, in m",
             [d["amp_outer"] for d in dict_list], label)
    fig_distr_funct("amplitude of outermost contour, in m",
                    [d["amp_outer"] for d in dict_list], label)

    fig_hist("amplitude of max-speed contour, in m",
             [d["amp_speed"] for d in dict_list], label)
    fig_distr_funct("amplitude of max-speed contour, in m",
                    [d["amp_speed"] for d in dict_list], label)
    
if __name__ == "__main__":
    import sys

    if len(sys.argv) != 2:
        sys.exit("Required argument: directory containing collection of "
                 "shapefiles")
        
    d = read(sys.argv[1])
    plot_all([d])
    plt.show()