piegy 2.1.0__cp38-cp38-win32.whl

piegy/tools/figure_tools.py

@@ -0,0 +1,238 @@
+'''
+Helper functions for making figures.
+
+Functions:
+- heatmap:          Make a heatmap based on input data. Sets title, text ... as well
+- bar:              Make a barplot. Sets title, text ... as well.
+- scatter:          Make a scatter plot. Sets title, text ... as well.
+- gen_title:        Generates a title when the plot is about an interval of time.
+- gen_text:         Generates a text about standard deviation info.
+- scale_interval:   scale interval if mod's data was already reduced.
+- ave_interval:     Calculates average value of data over a time interval.
+- ave_interval_1D:  Return in a 1D format.
+- config_mpl:       Configure Matplotlib parameters in a nice format
+'''
+
+
+import matplotlib.pyplot as plt
+import numpy as np
+import seaborn as sns
+
+# move ax a bit left if add text
+# default value is [0.125, 0.11, 0.9, 0.88]
+
+
+def heatmap(data, ax = None, cmap = "Greens", annot = False, fmt = '.3g', title = None, text = None):
+    '''
+    Helper function for making heatmaps.
+
+    Inputs:
+        data:   1D data for which you want to make a heatmap. 
+        ax:     matplotlib ax to plot on. 
+        cmap:   Color of heatmap. Uses matplotlib color maps
+        annot:  Whether to show numbers of every block.
+        fmt:    Number format for annotations. How many digits you want to keep.
+        title:  The title you want to add. None means no title.
+        text:   Adds some text in a text block at the top-right corner.
+
+    Returns:
+        fig:    Seaborn heatmap.
+    '''
+
+    if ax == None:
+        _, ax = plt.subplots()
+
+    if text != None:
+        ax.text(0.8, 1.025, text, size = 10, linespacing = 1.5, transform = ax.transAxes)
+
+    ax = sns.heatmap(data, ax = ax, cmap = cmap, annot = annot, fmt = fmt)
+    ax.set_title(title, x = 0.5, y = 1)
+    
+    return ax
+
+
+
+def bar(data, ax = None, color = "green", xlabel = None, ylabel = None, title = None, text = None):
+    '''
+    Helper Function for making barplots.
+
+    Inputs:
+        data:   2D data to make barplot.
+        ax:     matplotlib ax to plot on. 
+        color:  Uses Matplotlib colors.
+        xlabel, y_label: 
+                Label for axes.
+        title:  Title for the barplot.
+        text:   Adds some text in a text block at the top-right corner.
+    
+    Returns:
+        fig:    A Matplotlib barplot.
+    '''
+
+    N = np.array(data).size
+    xaxis = np.array([i for i in range(N)])
+    
+    if ax == None:
+        if N > 60:
+            # make figure larger if has more data points
+            _, ax = plt.subplots(figsize = (min(N * 0.107, 6.4), 4.8))
+        else:
+            _, ax = plt.subplots()
+
+    if text != None:
+        ax.text(0.7, 1.025, text, size = 10, linespacing = 1.5, transform = ax.transAxes)
+
+    ax.bar(x = xaxis, height = data, color = color)
+    ax.set_xlabel(xlabel)
+    ax.set_ylabel(ylabel)
+    ax.set_title(title, x = 0.5, y = 1)
+    
+    return ax
+
+
+
+def scatter(X, Y, ax = None, color = "orange", alpha = 0.25, xlabel = "x", ylabel = "y", title = None):
+    '''
+    Helper function for makeing scatter plots.
+
+    Inputs:
+        X:      x-coordinates of points.
+        Y:      y-coordinates of points.
+        ax:     matplotlib ax to plot on. 
+        Note color is Matplotlib colors.
+    
+    Returns:
+        fig:    A Matplotlib scatter plot.
+    '''
+    
+    if ax == None:
+        _, ax = plt.subplots(figsize = (7.2, 5.4))
+    ax.scatter(X, Y, color = color, alpha = alpha)
+    
+    ax.set_xlabel(xlabel)
+    ax.set_ylabel(ylabel)
+    ax.set_title(title)
+
+    return ax
+
+
+
+def gen_title(title, start, end):
+    '''
+    Generate a title for plot when it's about an interval of time.
+    '''
+    title += ", " + str(round(start * 100, 1)) + "~" + str(round(end * 100, 1)) + "%"
+    return title
+
+
+
+def gen_text(ave, std):
+    '''
+    Generate text about standard deviation info.
+    '''
+    text = "ave=" + str(round(ave, 2)) + ", std=" + str(round(std, 2))
+    return text
+
+
+
+def ave_interval(data, start_index, end_index):
+    '''
+    Calculate average value of data over an interval. Return a 2D np.array
+    Assume data is 3D with shape N x M x K, then takes average on the 3rd axis.
+
+    Input:
+        data:       3D np.array or list. Will take average on the 3rd axis.
+        start_index, end_index: 
+                    over what interval to take average.
+
+    Returns:
+        data_ave:   2D np.array with shape N x M, contains average value of data.
+    '''
+    
+    N = len(data)
+    M = len(data[0])
+    
+    # plot a particular record
+    if start_index == end_index:
+        start_index = end_index - 1
+        
+    data_ave = np.zeros((N, M))
+    
+    for i in range(N):
+        for j in range(M):
+            for k in range(start_index, end_index):
+                data_ave[i][j] += data[i][j][k]
+            data_ave[i][j] /= (end_index - start_index)
+    
+    return data_ave
+
+
+
+def ave_interval_1D(data, start_index, end_index):
+    '''
+    Calculate average value of data over an interval. Return a 1D np.array.
+    Assume data is 3D and has shape (1 x M x K) or (N x 1 x K). Then implicitly 'compress' that 1 and takes average on the 3rd axis.
+
+    Input:
+        data:       3D np.array or list. One of its dimensions must have size 1. Will take average on the 3rd axis.
+        start_index, end_index: 
+                    over what interval to take average.
+
+    Returns:
+        data_ave:   1D np.array with len N * M, contains average value of data.
+    '''
+    
+    N = len(data)
+    M = len(data[0])
+
+    if start_index == end_index:
+        start_index = end_index - 1
+        
+    data_ave = np.zeros(N * M)
+    
+    for i in range(N):
+        for j in range(M):
+            for k in range(start_index, end_index):
+                data_ave[i * M + j] += data[i][j][k]
+            data_ave[i * M + j] /= (end_index - start_index)
+    
+    return data_ave
+
+
+
+def scale_interval(interval, compress_itv):
+    # scale interval if mod's data was already reduced.
+    if compress_itv < 1:
+        raise ValueError('figures.scale_interval has compress_itv < 1:', compress_itv)
+
+    interval = int(interval / compress_itv)
+    if interval == 0:
+        print('Warning: data already smoothed by an interval: mod.compress_itv =', compress_itv, 'which is coarser than your', interval)
+        interval = 1
+
+    return interval
+
+
+
+def config_mpl(mpl):
+    '''
+    Configure Matplotlib figures
+    '''
+    mpl.rcParams['savefig.dpi'] = 300
+    
+    mpl.rcParams['font.family'] = 'serif'
+    mpl.rcParams['font.serif'] = plt.rcParams['font.serif']
+    mpl.rcParams['lines.linewidth'] = 1.75
+    mpl.rcParams['font.size'] = 11
+    mpl.rcParams['axes.labelsize'] = 12
+    mpl.rcParams['xtick.major.size'] = 10
+    mpl.rcParams['ytick.major.size'] = 9
+    mpl.rcParams['xtick.minor.size'] = 4
+    mpl.rcParams['ytick.minor.size'] = 4
+
+    mpl.rcParams['xtick.major.width'] = 1.5
+    mpl.rcParams['ytick.major.width'] = 1.5
+    mpl.rcParams['xtick.minor.width'] = 1.5
+    mpl.rcParams['ytick.minor.width'] = 1.5
+
+

piegy/tools/file_tools.py

@@ -0,0 +1,29 @@
+'''
+File-processing tools.
+
+Functions:
+- del_dirs:     Delete everything in a directory, as well as the directory itself.
+'''
+
+import os
+
+
+def del_dirs(dirs):
+    # Delete everything in a directory.
+    
+    subdirs_list = []
+    
+    for subdirs, dirs_, files in os.walk(dirs):
+        if subdirs not in subdirs_list:
+            subdirs_list.append(subdirs)
+            
+        for file in files:
+            path = os.path.join(subdirs, file)
+            if os.path.isfile(path):
+                os.remove(path)
+    
+    len_s = len(subdirs_list)
+    
+    for i in range(len_s):
+        os.rmdir(subdirs_list[len_s - i - 1])
+

piegy/videos.py

@@ -0,0 +1,303 @@
+'''
+Make mp4 videos for simulation results.
+
+Videos are made by:
+make every frame by figures.py functions, then put frames together into a video.
+
+Public Function:
+- make_video:   make video based simulation results.
+
+Private Functions
+- get_max_lim:  Get the max lim (interval) over many lims, and then expand it a bit for better accommodation.
+                Essentially takes union of those intervals. 
+- video_lim:    Find a large enough xlim and ylim for video.
+- make_mp4:     Put frames together into a mp4.
+others not documented here.
+
+'''
+
+
+from . import figures
+from .tools import file_tools as file_t
+from .tools import figure_tools as figure_t
+
+import matplotlib as mpl
+import matplotlib.pyplot as plt
+import numpy as np
+import os
+from moviepy import ImageSequenceClip
+
+
+# a list of supported figures
+SUPPORTED_FIGURES = ['UV_heatmap', 'pi_heatmap', 'UV_bar', 'pi_bar', 'UV_hist', 'pi_hist', 'UV_pi']
+
+
+# map function name to functios in figures.py
+# functions not in this dictionary is not supported for videos.
+FUNC_DICT = {'UV_heatmap': figures.UV_heatmap, 'UV_bar': figures.UV_bar, 'UV_hist': figures.UV_hist, 
+             'pi_heatmap': figures.pi_heatmap, 'pi_bar': figures.pi_bar, 'pi_hist': figures.pi_hist, 'UV_pi': figures.UV_pi}
+
+
+# Map some color maps to regular colors, used to change colors when an invalid color name is given
+SNS_PLT_COLOR_DICT = {'Greens': 'green', 'Purples': 'purple', 'BuPu': 'violet', 'YlGn': 'yellowgreen'}
+# Map regular colors to color maps
+PLT_SNS_COLOR_DICT = {'green': 'Greens', 'purple': 'Purples', 'violet': 'BuPu', 'yellowgreen': 'YlGn'}
+
+
+
+
+def convert_color(func_name, U_color, V_color):
+    '''
+    Converts some invalid colors.
+    If making heatmap videos but gave single colors, map to color maps.
+    If making barplot or histogram videos but gave single colors, map to Matplotlib
+    '''
+
+    if 'heatmap' in func_name:
+        # if making heatmaps but give regular colors
+        if U_color in PLT_SNS_COLOR_DICT.keys():
+            print('Making heatmaps, changed \'' + U_color + '\' to \'' + PLT_SNS_COLOR_DICT[U_color] + '\'')
+            U_color = PLT_SNS_COLOR_DICT[U_color]
+        if V_color in PLT_SNS_COLOR_DICT.keys():
+            print('Making heatmaps, changed \'' + V_color + '\' to \'' + PLT_SNS_COLOR_DICT[V_color] + '\'')
+            V_color = PLT_SNS_COLOR_DICT[V_color]
+        
+        return U_color, V_color
+
+    elif 'heatmap' not in func_name:
+        # if making barplots or histogram
+        if U_color in SNS_PLT_COLOR_DICT.keys():
+            print('Not making heatmaps, changed \'' + U_color + '\' to \'' + SNS_PLT_COLOR_DICT[U_color] + '\'')
+            U_color = SNS_PLT_COLOR_DICT[U_color]
+        if V_color in SNS_PLT_COLOR_DICT.keys():
+            print('Not making heatmaps, changed \'' + V_color + '\' to \'' + SNS_PLT_COLOR_DICT[V_color] + '\'')
+            V_color = SNS_PLT_COLOR_DICT[V_color]
+
+        return U_color, V_color
+
+
+
+def get_max_lim(lims):
+    '''
+    Get the max lim over many lims, i.e., the lowest lower bound and highest upper bound.
+    And then expand it a bit for better accommodation.
+
+    Input:
+        lim:    list or np.array, has form [lim1, lim2, ...]
+    
+    Returns:
+        A max lim which contains all lims.
+    '''
+
+    lims = np.array(lims)
+    
+    lim_min = np.min(lims[:, 0]) # min of min
+    lim_max = np.max(lims[:, 1]) # max of max
+    r = lim_max - lim_min
+
+    if lim_min != 0:
+        # negative values are reached
+        # extend both upper bound and lower bound 
+        return [lim_min - r * 0.05, lim_max + r * 0.05]
+    else:
+        # only extend upper bound
+        return [0, lim_max + r * 0.05]
+
+
+
+
+def frame_lim(mod, func, frames):
+    '''
+    Find a large enough xlim and ylim for frames, if not heatmaps.
+
+    Inputs:
+        mod:        A simulation.model object, the simulation results.
+        frames:     How many frame to make for the video.
+    
+    Returns:
+        xlim and ylim for U and V, 4 in total.
+    '''
+    
+    # take 10 samples and store their lims in list
+    U_xlist = []
+    U_ylist = []
+    V_xlist = []
+    V_ylist = []
+    
+    for i in range(10):
+        fig_U, ax_U = plt.subplots()
+        fig_V, ax_V = plt.subplots()
+        ax_U, ax_V = func(mod, ax_U = ax_U, ax_V = ax_V, start = i / 10, end = (i / 10 + 1 / frames))
+
+        U_xlist.append(ax_U.get_xlim())
+        U_ylist.append(ax_U.get_ylim())
+        V_xlist.append(ax_V.get_xlim())
+        V_ylist.append(ax_V.get_ylim())
+
+        plt.close(fig_U)
+        plt.close(fig_V)
+    
+    # get the largest 'range' based on the lists
+    U_xlim = get_max_lim(U_xlist)
+    U_ylim = get_max_lim(U_ylist)
+    V_xlim = get_max_lim(V_xlist)
+    V_ylim = get_max_lim(V_ylist)
+    
+    return U_xlim, U_ylim, V_xlim, V_ylim
+
+
+
+
+def frame_heatmap_lim(mod, func, frames):
+    '''
+    Find a large enough color bar lim for frames, if heatmaps.
+
+    Inputs:
+        mod:        A simulation.model object, the simulation results.
+        frames:     How many frame to make for the video.
+    
+    Returns:
+        clim for U and V
+    '''
+
+    U_list = []
+    V_list = []
+
+    for i in range(10):
+        fig_U, ax_U = plt.subplots()
+        fig_V, ax_V = plt.subplots()
+        ax_U, ax_V = func(mod, ax_U = ax_U, ax_V = ax_V, start = i / 10, end = (i / 10 + 1 / frames))
+
+        U_list.append(ax_U.collections[0].get_clim())
+        V_list.append(ax_V.collections[0].get_clim())
+
+        plt.close(fig_U)
+        plt.close(fig_V)
+
+    U_clim = get_max_lim(U_list)
+    V_clim = get_max_lim(V_list)
+
+    return U_clim, V_clim
+
+
+
+
+def make_mp4(video_dir, frame_dir, fps):
+    '''
+    Read .png from the frames folder and make into a mp4
+    Inputs:
+        video_dir:  where to save the video
+        frame_dirs: where to read frames from
+        fps:        frames per second
+    '''
+
+    frame_paths_incomplete = os.listdir(frame_dir)
+    frame_paths_incomplete.sort(key=lambda f: int(''.join(filter(str.isdigit, f))))
+    frame_paths = []
+    for path in frame_paths_incomplete:
+        if (path[-4:] == '.png') and ('frame' in path):
+            frame_paths.append(frame_dir + '/' + path)
+    clip = ImageSequenceClip(frame_paths, fps = fps) 
+    clip.write_videofile(video_dir, logger = None)
+
+
+
+
+def make_video(mod, func_name = 'UV_heatmap', frames = 100, dpi = 200, fps = 30, U_color = 'Greens', V_color = 'Purples', annot = False, fmt = '.3g', del_frames = False, dirs = 'videos'):
+    '''
+    Make a mp4 video based on simulation results.
+
+    Inputs:
+    - mod:            a simulation.model object, the simulation results.
+    - func_name:      what function to use to make the frames. Should be one of the functions in figures.py
+    - frames:         how many frames to make. Use more frames for more smooth evolutions.
+    - dpi:            dots per inch.
+    - fps:            frames per second.
+    - U_color:        color for U's videos. Color maps or regular colors, based on what function you use.
+    - V_color:        color for V's videos.
+    - annot:          used by heatmaps. Whether to show numbers.
+    - fmt:            number format
+    - del_frames:     whether to delete frames after making video.
+    - dirs:           where to store the frames and videos.
+    '''
+    
+    if func_name not in FUNC_DICT.keys():
+        raise ValueError(func_name + ' not supported for videos.')
+    func = FUNC_DICT[func_name]
+
+    # convert color if invalid colors are given
+    U_color, V_color = convert_color(func_name, U_color, V_color)
+    
+    # print progress
+    one_progress = frames / 100
+    current_progress = one_progress
+    
+    if 'heatmap' in func_name:
+        # make sure a fixed color bar for all frames
+        U_clim, V_clim = frame_heatmap_lim(mod, func, frames)
+    else:
+        # make sure y axis not changing if not making heatmaps
+        U_xlim, U_ylim, V_xlim, V_ylim = frame_lim(mod, func, frames)
+
+    
+    U_frame_dirs = dirs + '/U-' + func_name
+    V_frame_dirs = dirs + '/V-' + func_name
+    
+    if os.path.exists(U_frame_dirs):
+        file_t.del_dirs(U_frame_dirs)
+    os.makedirs(U_frame_dirs)
+    if os.path.exists(V_frame_dirs):
+        file_t.del_dirs(V_frame_dirs)
+    os.makedirs(V_frame_dirs)
+        
+    #### for loop ####
+    
+    for i in range(frames):
+        if i > current_progress:
+            print('making frames', round(i / frames * 100), '%', end = '\r')
+            current_progress += one_progress
+        
+        fig_U, ax_U = plt.subplots(figsize = (6.4, 4.8))#, constrained_layout = True)
+        fig_V, ax_V = plt.subplots(figsize = (6.4, 4.8))#, constrained_layout = True)
+
+        if 'heatmap' in func_name:
+            ax_U, ax_V = func(mod, ax_U = ax_U, ax_V = ax_V, U_color = U_color, V_color = V_color, start = i / frames, end = (i + 1) / frames, annot = annot, fmt = fmt)
+        else:
+            ax_U, ax_V = func(mod, ax_U = ax_U, ax_V = ax_V, U_color = U_color, V_color = V_color, start = i / frames, end = (i + 1) / frames)
+        
+        if 'heatmap' in func_name:
+            ax_U.collections[0].set_clim(U_clim)
+            ax_V.collections[0].set_clim(V_clim)
+        else:
+            # make sure y axis not changing if not heatmap and not UV_pi
+            ax_U.set_ylim(U_ylim)
+            ax_V.set_ylim(V_ylim)
+            if ('hist' in func_name) or (func_name == 'UV_pi'):
+                # need to set xlim as well for UV_pi and histograms
+                ax_V.set_xlim(U_xlim)
+                ax_V.set_xlim(V_xlim)
+
+        fig_U.savefig(U_frame_dirs + '/' + 'U_frame_' + str(i) + '.png', pad_inches = 0.25, dpi = dpi)
+        fig_V.savefig(V_frame_dirs + '/' + 'V_frame_' + str(i) + '.png', pad_inches = 0.25, dpi = dpi)
+        
+        plt.close(fig_U)
+        plt.close(fig_V)
+        
+    #### for loop ends ####
+    
+    # frames done
+    print('making mp4...      ', end = '\r')
+    
+    # make videos based on frames
+    make_mp4(dirs + '/U-' + func_name + '.mp4', U_frame_dirs, fps)
+    make_mp4(dirs + '/V-' + func_name + '.mp4', V_frame_dirs, fps)
+    
+    if del_frames:
+        file_t.del_dirs(U_frame_dirs)
+        file_t.del_dirs(V_frame_dirs)
+        print('video saved: ' + dirs + ', frames deleted')
+    else:
+        print('video saved: ' + dirs + '      ')
+
+
+

piegy-2.1.0.dist-info/LICENSE.txt

@@ -0,0 +1,28 @@
+BSD 3-Clause License
+
+Copyright (c) 2025, Chenning Xu
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this
+   list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice,
+   this list of conditions and the following disclaimer in the documentation
+   and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its
+   contributors may be used to endorse or promote products derived from
+   this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.