[levels_to_step_plot] #python #plot #step_plot

levels_to_step_plot.py

import numpy as np

def levels_to_steps_post(x_vec, y_vec):
    """Provides data like matplotlib.pyplot.step with where set to 'post'"""
    
    x_new = [None] * (len(x_vec) * 2 - 1)
    y_new = [None] * (len(x_vec) * 2 - 1)
    
    for i in range(len(x_vec)-1):
        x_new[2*i] = x_vec[i]
        y_new[2*i] = y_vec[i]
        
        x_new[2*i + 1] = x_vec[i+1]
        y_new[2*i + 1] = y_vec[i]
        
    x_new[len(x_vec)*2-2] = x_vec[len(x_vec)-1]
    y_new[len(x_vec)*2-2] = y_vec[len(x_vec)-1]
    
    return np.array(x_new), np.array(y_new)

  
def levels_to_steps(level_list):
    """
    Convert list of 2-element tuples that contain duration and level data (x, y) to 1d x,y numpy arrays. 
    The arrays plotted give a step plot. 
    
    Args:
        level_list (list of tuples with two number elements): [(x1, y1), (x2, y2), ...]
        
    Returns:
        numpy.ndarray: 1d
        numpy.ndarray: 1d
        
    Example:
        >> x, y = levels_to_steps([(5, 4), (5, 8), (10, 5)])
        
        >> x
        array([0,  0,  5,  5, 10, 10, 20, 20])
        
        >> y
        array([0,  4,  4,  8,  8,  5,  5,  0])
        
    
    """
    x_list = [0] * (len(level_list)* 2 + 2)
    y_list = [0] * (len(level_list)* 2 + 2)
    
    cumulative_time = 0
    for i in range(len(level_list)):
        cumulative_time += level_list[i][0]
        for j in range(2):
            x_list[2 + i*2 + j] = cumulative_time
            
    for i in range(len(level_list)):
        y = level_list[i][1]
        for j in range(2):
            y_list[1 + i*2 + j] = y
    
    return np.array(x_list), np.array(y_list)