piegy 1.0.0__py3-none-any.whl → 1.1.1__py3-none-any.whl

piegy/__version__.py

@@ -13,5 +13,7 @@ version history:
 0.1.9: first round of full debugging
 
 1.0.0: first version in PyPI
+1.1.0: debugging. Updated a range of functions, in the following modules: figures, videos, test_var, model, figure_tools
+1.1.1: minor debugging in model module.
 
 '''

piegy/data_tools.py

@@ -118,12 +118,14 @@ def read_data(dirs):
 
     # outputs
     try:
-        sim.set_data(False, data[2][0], data[2][1], data[2][2], data[2][3], data[2][4], data[2][5])
+        sim.set_data(data_empty = False, max_record = data[2][0], compress_itv = data[2][1], 
+                    U = data[2][2], V = data[2][3], U_pi = data[2][4], V_pi = data[2][5])
     except:
         raise ValueError('Invalid simulation results saved in data')
     
     return sim
-    
+
+
 
 
 

piegy/figures.py

@@ -37,20 +37,20 @@ import matplotlib.pyplot as plt
 import numpy as np
 
 
-
 # curve type in plot
 # used by UV_dyna, UV_std, and pi_dyna
 CURVE_TYPE = '-'
 
 
 
-def UV_heatmap(sim, U_color = 'Purples', V_color = 'Greens', start = 0.95, end = 1.0, annot = False, fmt = '.3g'):
+def UV_heatmap(sim, ax_U = None, ax_V = None, U_color = 'Purples', V_color = 'Greens', start = 0.95, end = 1.0, annot = False, fmt = '.3g'):
     '''
     Makes two heatmaps for U, V average distribution over a time interval, respectively. Works best for 2D space.
     1D works as well, but figures look bad.
 
     Inputs:
         sim:        A model.simulation object.
+        ax_U, ax_V: matplotlib axes to plot on. New axes will be created if None is given.
         U_color:    Color for U's heatmap, uses matplotlib color maps.
         V_color:    Color for V's heatmap.
         start:      (0,1) float, where the interval should start from. Intended as a 'percentage'. 
@@ -60,7 +60,7 @@ def UV_heatmap(sim, U_color = 'Purples', V_color = 'Greens', start = 0.95, end =
         fmt:        Number format for annotations. How many digits you want to keep. Please set annot = True first and then use fmt.
 
     Returns:
-        fig1, fig2: Two heatmaps of U, V distribution.
+        ax_U, ax_V: matplotlib axes with heatmaps of U, V distribution plotted upon.
     '''
     
     start_index = int(start * sim.max_record)
@@ -77,20 +77,21 @@ def UV_heatmap(sim, U_color = 'Purples', V_color = 'Greens', start = 0.95, end =
     V_title = figure_t.gen_title('V', start, end)
     V_text = figure_t.gen_text(np.mean(V_ave), np.std(V_ave))
 
-    fig1 = figure_t.heatmap(U_ave, U_color, annot, fmt, U_title, U_text)
-    fig2 = figure_t.heatmap(V_ave, V_color, annot, fmt, V_title, V_text)
+    ax_U = figure_t.heatmap(U_ave, ax_U, U_color, annot, fmt, U_title, U_text)
+    ax_V = figure_t.heatmap(V_ave, ax_V, V_color, annot, fmt, V_title, V_text)
         
-    return fig1, fig2
+    return ax_U, ax_V
     
 
 
-def UV_bar(sim, U_color = 'purple', V_color = 'green', start = 0.95, end = 1.0):
+def UV_bar(sim, ax_U = None, ax_V = None, U_color = 'purple', V_color = 'green', start = 0.95, end = 1.0):
     '''
     Makes two barplots for U, V average distribution over a time interval. Works best for 1D space.
     2D works as well, but figures look bad.
 
     Inputs:
         sim:        A model.simulation object.
+        ax_U, ax_V: matplotlib axes to plot on. New axes will be created if None is given.
         U_color:    Color of U's barplot. Uses Matplotlib colors.
                     See available colors at: https://matplotlib.org/stable/gallery/color/named_colors.html
         V_color:    Color of V's barplot. Uses Matplotlib colors.
@@ -98,7 +99,7 @@ def UV_bar(sim, U_color = 'purple', V_color = 'green', start = 0.95, end = 1.0):
         end:        (0,1) float. Where you want the interval to end.
 
     Returns:
-        fig1, fig2: Two Matplotlib bar plots, for U and V, respectively.
+        ax_U, ax_V: matplotlib axes with bar plots for U and V plotted upon.
     '''
     
     start_index = int(start * sim.max_record)
@@ -114,15 +115,15 @@ def UV_bar(sim, U_color = 'purple', V_color = 'green', start = 0.95, end = 1.0):
     V_title = figure_t.gen_title('V', start, end)
     V_text = figure_t.gen_text(np.mean(V_ave), np.std(V_ave))
 
-    fig1 = figure_t.bar(U_ave, color = U_color, xlabel = 'patches', ylabel = 'U', title = U_title, text = U_text)
-    fig2 = figure_t.bar(V_ave, color = V_color, xlabel = 'patches', ylabel = 'V', title = V_title, text = V_text)
+    ax_U = figure_t.bar(U_ave, ax = ax_U, color = U_color, xlabel = 'patches', ylabel = 'U', title = U_title, text = U_text)
+    ax_V = figure_t.bar(V_ave, ax = ax_V, color = V_color, xlabel = 'patches', ylabel = 'V', title = V_title, text = V_text)
 
-    return fig1, fig2
+    return ax_U, ax_V
 
 
 
 
-def UV_dyna(sim, interval = 20, grid = True):
+def UV_dyna(sim, ax = None, interval = 20, grid = True):
     '''
     Plots how total U, V change overtime.
     The curves are not directly based on every single data point. 
@@ -132,13 +133,14 @@ def UV_dyna(sim, interval = 20, grid = True):
 
     Inputs:
         sim:        A model.simulation object.
+        ax:         matplotlib ax to plot on. New ax will be created if None is given.
         interval:   How many data points to take average over. Larger value makes curves smoother, but also loses local fluctuations.
                     NOTE: this interval doesn't overlap with sim.compress_itv. 
                     e.g. you already took average over every 20 data points, then using interval <= 20 here has no smoothing effect.
         grid:       Whether to add grid lines to plot.
     
     Returns:
-        fig:        A Matplotlib figure, contains U's, V's, and U+V population.
+        ax:        matplotlib ax, contains U's, V's, and sum of U & V population.
     '''
     
     # store the average values in lists
@@ -160,26 +162,29 @@ def UV_dyna(sim, interval = 20, grid = True):
     #### plot ####   
     xaxis = np.linspace(0, sim.maxtime, len(U_curve))
 
-    fig, ax = plt.subplots()
+    if ax == None:
+        _, ax = plt.subplots()
     ax.grid(grid)
     ax.plot(xaxis, U_curve, CURVE_TYPE, label = 'U')
     ax.plot(xaxis, V_curve, CURVE_TYPE, label = 'V')
     ax.plot(xaxis, total_curve, CURVE_TYPE, label = 'total')
-    ax.title.set_text('U & V over time')
+    ax.set_xlabel('Time')
+    ax.set_ylabel('Population')
+    ax.set_title('Population U & V Over Time')
     ax.legend()
 
-    return fig
+    return ax
 
 
 
 
-def UV_hist(sim, U_color = 'purple', V_color = 'green', start = 0.95, end = 1.0):
+def UV_hist(sim, ax_U = None, ax_V = None, U_color = 'purple', V_color = 'green', start = 0.95, end = 1.0):
     '''
     Makes density histograms for U, V's average distribution over an interval.
     Sometimes it may not be shown in density plots due to matplotlib features.
 
     Returns:
-        fig1, fig2: Two Matplotlib histograms, for U and V, respectively.
+        ax_U, ax_V: matplotlib axes with heatmaps of U, V population density plotted upon.
     '''
 
     start_index = int(start * sim.max_record)
@@ -190,30 +195,32 @@ def UV_hist(sim, U_color = 'purple', V_color = 'green', start = 0.95, end = 1.0)
     
     #### plot ####
     
-    fig1, ax1 = plt.subplots()
-    ax1.set_xlabel('U')
-    ax1.set_ylabel('density')
-    ax1.hist(U_ave, color = U_color, density = True)
-    ax1.title.set_text(figure_t.gen_title('U hist', start, end))
+    if ax_U == None:
+        _, ax_U = plt.subplots()
+    ax_U.set_xlabel('Population U')
+    ax_U.set_ylabel('Density')
+    ax_U.hist(U_ave, color = U_color, density = True)
+    ax_U.set_title(figure_t.gen_title('U Hist', start, end))
     
-    fig2, ax2 = plt.subplots()
-    ax2.set_xlabel('V')
-    ax2.set_ylabel('density')
-    ax2.hist(V_ave, color = V_color, density = True)
-    ax2.title.set_text(figure_t.gen_title('V hist', start, end))
+    if ax_V == None:
+        _, ax_V = plt.subplots()
+    ax_V.set_xlabel('Population V')
+    ax_V.set_ylabel('Density')
+    ax_V.hist(V_ave, color = V_color, density = True)
+    ax_V.set_title(figure_t.gen_title('V Hist', start, end))
 
-    return fig1, fig2
+    return ax_U, ax_V
 
 
 
 
-def UV_std(sim, interval = 20, grid = True):
+def UV_std(sim, ax = None, interval = 20, grid = True):
     '''
     Plots how standard deviation of U, V change over time.
     Takes average over many small interval to smooth out local fluctuations.
 
     Returns:
-        fig:    A Matplotlib figure, contains U's and V's std curves.
+        ax:    matplotlib ax, contains U's and V's std curves.
     '''
 
     interval = figure_t.scale_interval(interval, sim.compress_itv)
@@ -232,15 +239,17 @@ def UV_std(sim, interval = 20, grid = True):
     #### plot ####
     xaxis = np.linspace(0, sim.maxtime, len(U_std))
 
-    fig, ax = plt.subplots()
+    if ax == None:
+        _, ax = plt.subplots()
     ax.grid(grid)
     ax.plot(xaxis, U_std, CURVE_TYPE, label = 'U std')
     ax.plot(xaxis, V_std, CURVE_TYPE, label = 'V std')
     ax.legend()
-    ax.set_xlabel('time (records)', fontsize = 11)
-    ax.title.set_text('std_dev over time')
-    
-    return fig
+    ax.set_xlabel('Time')
+    ax.set_ylabel('Std Dev')
+    ax.set_title('Population Std-Dev Dynamics')
+
+    return ax
 
 
 
@@ -270,7 +279,7 @@ def UV_expected_val(sim):
 
 
 
-def UV_expected(sim, U_color = 'Purples', V_color = 'Greens', annot = False, fmt = '.3g'):
+def UV_expected(sim, ax_U = None, ax_V = None, U_color = 'Purples', V_color = 'Greens', annot = False, fmt = '.3g'):
     '''
     Calculate expected population distribution based on matrices, assuming no migration.
     For the formulas, see stochastic_mode.expected_UV
@@ -279,8 +288,8 @@ def UV_expected(sim, U_color = 'Purples', V_color = 'Greens', annot = False, fmt
         Note the colors are color maps.
     
     Returns:
-    fig1, fig2: If 2D (N and M both > 1), then fig1 and fig2 are heatmaps.
-                If 1D (N or M == 1), then fig1 and fig2 are barplots.
+    ax_U, ax_V: If 2D (N and M both > 1), then ax_U and ax_V are heatmaps.
+                If 1D (N or M == 1), then ax_U and ax_V are barplots.
     '''
     
     U_expected, V_expected = UV_expected_val(sim)
@@ -292,20 +301,20 @@ def UV_expected(sim, U_color = 'Purples', V_color = 'Greens', annot = False, fmt
     
     if (sim.N != 1) and (sim.M != 1):
         # 2D
-        fig1 = figure_t.heatmap(U_expected, U_color, annot, fmt, title = 'Expected U', text = U_text)
-        fig2 = figure_t.heatmap(V_expected, V_color, annot, fmt, title = 'Expected V', text = V_text)
+        ax_U = figure_t.heatmap(U_expected, ax_U, U_color, annot, fmt, title = 'Expected U', text = U_text)
+        ax_V = figure_t.heatmap(V_expected, ax_V, V_color, annot, fmt, title = 'Expected V', text = V_text)
 
     else:
         # 1D     
-        fig1 = figure_t.bar(U_expected.flatten(), color = U_color, xlabel = 'patches', ylabel = 'popu', title = 'Expected U', text = U_text)
-        fig2 = figure_t.bar(V_expected.flatten(), color = V_color, xlabel = 'patches', ylabel = 'popu', title = 'Expected V', text = V_text)
+        ax_U = figure_t.bar(U_expected.flatten(), ax_U, color = U_color, xlabel = 'patches', ylabel = 'popu', title = 'Expected Population U', text = U_text)
+        ax_V = figure_t.bar(V_expected.flatten(), ax_V, color = V_color, xlabel = 'patches', ylabel = 'popu', title = 'Expected Population V', text = V_text)
 
-    return fig1, fig2
+    return ax_U, ax_V
 
 
 
 
-def pi_heatmap(sim, U_color = 'BuPu', V_color = 'YlGn', start = 0.95, end = 1.0, annot = False, fmt = '.3g'):
+def pi_heatmap(sim, ax_U = None, ax_V = None, U_color = 'BuPu', V_color = 'YlGn', start = 0.95, end = 1.0, annot = False, fmt = '.3g'):
     '''
     Make heatmaps for payoff in a specified interval.
     Works best for 2D. 1D works as well, but figures look bad.
@@ -314,7 +323,7 @@ def pi_heatmap(sim, U_color = 'BuPu', V_color = 'YlGn', start = 0.95, end = 1.0,
         Note the colors are matplotlib color maps.
 
     Returns:
-        fig1, fig2: Seaborn heatmaps, for U's & V's payoff distribution, respectively.
+        ax_U, ax_V: Seaborn heatmaps, for U's & V's payoff distribution, respectively.
     '''
     
     start_index = int(sim.max_record * start)
@@ -323,26 +332,26 @@ def pi_heatmap(sim, U_color = 'BuPu', V_color = 'YlGn', start = 0.95, end = 1.0,
     U_pi_ave = figure_t.ave_interval(sim.U_pi, start_index, end_index)
     V_pi_ave = figure_t.ave_interval(sim.V_pi, start_index, end_index)
     
-    U_title = figure_t.gen_title('U_pi', start, end)
+    U_title = figure_t.gen_title('Payoff ' + r'$p_U$', start, end)
     U_text = figure_t.gen_text(np.mean(U_pi_ave), np.std(U_pi_ave))
-    V_title = figure_t.gen_title('V_pi', start, end)
+    V_title = figure_t.gen_title('Payoff ' + r'$p_V$', start, end)
     V_text = figure_t.gen_text(np.mean(V_pi_ave), np.std(V_pi_ave))
     
-    fig1 = figure_t.heatmap(U_pi_ave, U_color, annot, fmt, U_title, U_text)
-    fig2 = figure_t.heatmap(V_pi_ave, V_color, annot, fmt, V_title, V_text)
+    ax_U = figure_t.heatmap(U_pi_ave, ax_U, U_color, annot, fmt, U_title, U_text)
+    ax_V = figure_t.heatmap(V_pi_ave, ax_V, V_color, annot, fmt, V_title, V_text)
 
-    return fig1, fig2
+    return ax_U, ax_V
 
 
 
 
-def pi_bar(sim, U_color = 'violet', V_color = 'yellowgreen', start = 0.95, end = 1.0):
+def pi_bar(sim, ax_U = None, ax_V = None, U_color = 'violet', V_color = 'yellowgreen', start = 0.95, end = 1.0):
     '''
     Make barplot for payoff in a specified interval.
     Works best for 1D. 2D works as well, but figures look bad.
 
     Returns:
-        fig1, fig2: Matplotlib barplots, for U's and V's payoff distribution, respectively.
+        ax_U, ax_V: matplotlib axes with barplots of U and V payoff distribution plotted upon.
     '''
     
     start_index = int(sim.max_record * start)
@@ -351,25 +360,25 @@ def pi_bar(sim, U_color = 'violet', V_color = 'yellowgreen', start = 0.95, end =
     U_pi_ave = figure_t.ave_interval_1D(sim.U_pi, start_index, end_index)
     V_pi_ave = figure_t.ave_interval_1D(sim.V_pi, start_index, end_index)
     
-    U_title = figure_t.gen_title('U_pi', start, end)
+    U_title = figure_t.gen_title(r'$p_U$', start, end)
     U_text = figure_t.gen_text(np.mean(U_pi_ave), np.std(U_pi_ave))
-    V_title = figure_t.gen_title('V_pi', start, end)
+    V_title = figure_t.gen_title(r'$p_V$', start, end)
     V_text = figure_t.gen_text(np.mean(V_pi_ave), np.std(V_pi_ave))
     
-    fig1 = figure_t.bar(U_pi_ave, U_color, 'patches', 'pi', U_title, U_text)
-    fig2 = figure_t.bar(V_pi_ave, V_color, 'patches', 'pi', V_title, V_text)
+    ax_U = figure_t.bar(U_pi_ave, ax_U, U_color, 'Patches', 'Payoff ' + r'$p_U$', U_title, U_text)
+    ax_V = figure_t.bar(V_pi_ave, ax_V, V_color, 'Patches', 'Payoff ' + r'$p_V$', V_title, V_text)
 
-    return fig1, fig2
+    return ax_U, ax_V
 
 
 
 
-def pi_dyna(sim, interval = 20, grid = True):
+def pi_dyna(sim, ax = None, interval = 20, grid = True):
     '''
     Plot how payoffs change over time.
 
     Returns:
-        fig:    Matplotlib figure of U's, V's, and U+V payoff, either total or not.
+        ax:    matplotlib ax of U's, V's, and sum of U & V payoff.
     '''
     
     U_curve = []
@@ -390,27 +399,30 @@ def pi_dyna(sim, interval = 20, grid = True):
     #### plot ####    
     xaxis = np.linspace(0, sim.maxtime, len(U_curve))
     
-    fig, ax = plt.subplots()
+    if ax == None:
+        _, ax = plt.subplots()
     ax.grid(grid)
-    ax.plot(xaxis, U_curve, CURVE_TYPE, label = 'U_pi')
-    ax.plot(xaxis, V_curve, CURVE_TYPE, label = 'V_pi')
+    ax.plot(xaxis, U_curve, CURVE_TYPE, label = r'$p_U$')
+    ax.plot(xaxis, V_curve, CURVE_TYPE, label = r'$p_V$')
     ax.plot(xaxis, total_curve, CURVE_TYPE, label = 'total')
     ax.set_xlim(0, sim.maxtime)
-    ax.title.set_text('U&V _pi over time')
+    ax.set_xlabel('Time')
+    ax.set_ylabel('Payoff')
+    ax.set_title('Payoff ' + r'$p_U$' + ' & ' + r'$p_V$' + ' over time')
     ax.legend()
 
-    return fig
+    return ax
 
 
 
 
-def pi_hist(sim, U_color = 'violet', V_color = 'yellowgreen', start = 0.95, end = 1.0):
+def pi_hist(sim, ax_U = None, ax_V = None, U_color = 'violet', V_color = 'yellowgreen', start = 0.95, end = 1.0):
     '''
     Makes deensity histograms of U's and V's payoffs in a sepcified interval.
     Sometimes it may not be shown in density plots due to matplotlib features.
     
     Returns:
-        fig1, fig2:     histogram of U's and V's payoff.
+        ax_U, ax_V:     histogram of U's and V's payoff.
     '''
 
     start_index = int(start * sim.max_record)
@@ -421,29 +433,31 @@ def pi_hist(sim, U_color = 'violet', V_color = 'yellowgreen', start = 0.95, end
     
     #### plot ####
     
-    fig1, ax1 = plt.subplots()
-    ax1.set_xlabel('U_pi')
-    ax1.set_ylabel('density')
-    ax1.hist(U_pi_ave, color = U_color, density = True)
-    ax1.title.set_text(figure_t.gen_title('U_pi hist', start, end))
+    if ax_U == None:
+        _, ax_U = plt.subplots()
+    ax_U.set_xlabel('Payoff ' + r'$p_U$')
+    ax_U.set_ylabel('Density')
+    ax_U.hist(U_pi_ave, color = U_color, density = True)
+    ax_U.set_title(figure_t.gen_title('Payoff ' + r'$p_U$' + ' Hist', start, end))
     
-    fig2, ax2 = plt.subplots()
-    ax2.set_xlabel('V_pi')
-    ax2.set_ylabel('density')
-    ax2.hist(V_pi_ave, color = V_color, density = True)
-    ax2.title.set_text(figure_t.gen_title('V_pi hist', start, end))
+    if ax_V == None:
+        _, ax_V = plt.subplots()
+    ax_V.set_xlabel('Payoff ' + r'$p_V$')
+    ax_V.set_ylabel('Density')
+    ax_V.hist(V_pi_ave, color = V_color, density = True)
+    ax_V.set_title(figure_t.gen_title('Payoff ' + r'$p_V$' + ' Hist', start, end))
 
-    return fig1, fig2
+    return ax_U, ax_V
 
 
 
 
-def pi_std(sim, interval = 20, grid = True):
+def pi_std(sim, ax = None, interval = 20, grid = True):
     '''
     Plots how standard deviation of payoff change over time.
 
     Returns:
-        fig:    Matplotlib figure of the std of payoffs.
+        ax:    matplotlib ax of the std of payoffs.
     '''
     
     
@@ -463,26 +477,28 @@ def pi_std(sim, interval = 20, grid = True):
     #### plot ####
     xaxis = np.linspace(0, sim.maxtime, len(U_pi_std))
     
-    fig, ax = plt.subplots()
+    if ax == None:
+        _, ax = plt.subplots()
     ax.grid(grid)
-    ax.plot(xaxis, U_pi_std, CURVE_TYPE, label = 'U_pi std')
-    ax.plot(xaxis, V_pi_std, CURVE_TYPE, label = 'V_pi std')
+    ax.plot(xaxis, U_pi_std, CURVE_TYPE, label = r'$p_U$' + ' std')
+    ax.plot(xaxis, V_pi_std, CURVE_TYPE, label = r'$p_V$' + ' std')
     ax.legend()
-    ax.set_xlabel('time (records)', fontsize = 11)
-    ax.title.set_text('std over time')
+    ax.set_xlabel('Time')
+    ax.set_ylabel('Std Dev')
+    ax.set_title('Payoff Std-Dev Dynamics')
     
-    return fig
+    return ax
 
 
 
 
-def UV_pi(sim, U_color = 'violet', V_color = 'yellowgreen', alpha = 0.25, start = 0.95, end = 1.0):
+def UV_pi(sim, ax_U = None, ax_V = None, U_color = 'violet', V_color = 'yellowgreen', alpha = 0.4, start = 0.95, end = 1.0):
     '''
     Make two scatter plots: x-axes are population and y-axes are payoff in a specified time interval.
     Reveals relationship between population and payoff.
 
     Returns:
-        fig1, fig2: U's and V's population-payoff scatter plots.
+        ax_U, ax_V: matplotlib axes with U and V population-payoff scatter plots.
     '''
     
     start_index = int(start * sim.max_record)
@@ -495,9 +511,9 @@ def UV_pi(sim, U_color = 'violet', V_color = 'yellowgreen', alpha = 0.25, start
     V_pi_ave = figure_t.ave_interval(sim.V_pi, start_index, end_index)
     
     
-    fig1 = figure_t.scatter(U_ave, U_pi_ave, U_color, alpha, xlabel = 'U', ylabel = 'U_pi', title = 'U - U_pi')
-    fig2 = figure_t.scatter(V_ave, V_pi_ave, V_color, alpha, xlabel = 'V', ylabel = 'V_pi', title = 'V - V_pi')
+    ax_U = figure_t.scatter(U_ave, U_pi_ave, ax_U, U_color, alpha, xlabel = 'U', ylabel = 'Payoff ' + r'$p_U$', title = 'U - ' + r'$p_U$')
+    ax_V = figure_t.scatter(V_ave, V_pi_ave, ax_V, V_color, alpha, xlabel = 'V', ylabel = 'Payoff ' + r'$p_V$', title = 'V - ' + r'$p_V$')
     
-    return fig1, fig2
+    return ax_U, ax_V
 
 

piegy/model.py

@@ -41,9 +41,6 @@ PI_DTYPE = 'float64'
 # data type for storing rates in single_test an single_init
 RATES_DTYPE = 'float64'
 
-# store e locally, slightly increases speed
-MATH_E = math.e
-
 
 class patch:
     '''
@@ -176,8 +173,8 @@ class patch:
 
         for i in range(4):
             if self.nb[i] != None:
-                U_weight[i] = 1 + pow(MATH_E, self.w1 * self.nb[i].U_pi)  # use your own functions!
-                V_weight[i] = 1 + pow(MATH_E, self.w2 * self.nb[i].V_pi)
+                U_weight[i] = 1 + pow(math.e, self.w1 * self.nb[i].U_pi)
+                V_weight[i] = 1 + pow(math.e, self.w2 * self.nb[i].V_pi)
 
         mu1_U = self.mu1 * self.U
         mu2_V = self.mu2 * self.V
@@ -358,14 +355,11 @@ class simulation:
         if print_pct <= 0:
             raise ValueError('Please use an int > 0 for print_pct or None for not printing progress.')
         
-        if type(seed) != int:
-            raise TypeError('Please use an int as seed')
-        if seed < 0:
-            raise ValueError('Please use a non-negative int as seed.')
+        # seed is checked by number
         
 
     def check_valid_data(self, data_empty, max_record, compress_itv, U, V, U_pi, V_pi):
-        # check whether a set of data is valid
+        # check whether a set of data is valid, used when reading a saved model
         if type(data_empty) != bool:
             raise TypeError('data_empty not a bool')
         

piegy/test_var.py

@@ -48,6 +48,10 @@ DOTTED_CURVE_TYPE = 'o-'
 # map patch_var name to index in the patch class (in stochastic_model.py)
 PATCH_VAR_DICT = {'mu1': 0, 'mu2': 1, 'w1': 2, 'w2': 3, 'kappa1': 4, 'kappa2': 5}
 
+# display name (in latex)
+PATCH_VAR_DISP = {'mu1': r'$\mu_U$', 'mu2': r'$\mu_V$', 'w1': r'$w_U$', 'w2': r'$w_V$', 
+                  'kappa1': r'$\kappa_U$', 'kappa2': r'$\kappa_V$'}
+
 
 
 def test_var1(sim, var, values, dirs, compress_itv = None, predict_runtime = False):
@@ -144,7 +148,7 @@ def test_var2(sim, var1, var2, values1, values2, dirs, compress_itv = None, pred
 
 
 
-def var_UV1(var, values, var_dirs, start = 0.95, end = 1.0, U_color = 'purple', V_color = 'green'):
+def var_UV1(var, values, var_dirs, ax_U = None, ax_V = None, start = 0.95, end = 1.0, U_color = 'purple', V_color = 'green'):
     '''
     Plot function for test_var1, plot how var influences U, V population.
     Make U, V - var curve in two figures, with y-axis being total population at the end of simulations, x-axis being var's values.
@@ -153,12 +157,13 @@ def var_UV1(var, values, var_dirs, start = 0.95, end = 1.0, U_color = 'purple',
     - var:        str, which variable was tested.
     - values:     1D list or np.array, which values were tested.
     - var_dirs:   return value of test_var1, a 1D list of directories where all data are stored.
+    - ax_U, ax_V:   matplotlib axes to plot on.
     - start, end: floats, give an interval of time over which to take average and make plot.
                 For example, start = 0.95, end = 1.0 are to take average over the last 5% time of results; 
                 essentially plots how var influences equilibrium population.
     
     Returns:
-    - fig1, fig2: two figures for U, V, respectively.
+    - ax_U, ax_V: axes for U, V, respectively.
     '''
 
     # average value of U, V over the interval. One entry for one value of var
@@ -181,25 +186,28 @@ def var_UV1(var, values, var_dirs, start = 0.95, end = 1.0, U_color = 'purple',
         V_ave.append(sum(figure_t.ave_interval_1D(simk.V, start_index, end_index)) / NM)
         
     #### plot ####
+    if ax_U == None:
+        _, ax_U = plt.subplots()
+    if ax_V == None:
+        _, ax_V = plt.subplots()
+    var_disp = PATCH_VAR_DISP[var]
     
-    fig1, ax1 = plt.subplots()
-    ax1.set_xlabel(var)
-    ax1.set_ylabel('U')
-    ax1.plot(values, U_ave, DOTTED_CURVE_TYPE, color = U_color)
-    ax1.title.set_text(figure_t.gen_title(var + ' - U', start, end))
+    ax_U.set_xlabel(var_disp)
+    ax_U.set_ylabel(r'$U$')
+    ax_U.plot(values, U_ave, DOTTED_CURVE_TYPE, color = U_color)
+    ax_U.set_title(figure_t.gen_title(var_disp + r'$\,-\,U$', start, end))
     
-    fig2, ax2 = plt.subplots()
-    ax2.set_xlabel(var)
-    ax2.set_ylabel('V')
-    ax2.plot(values, V_ave, DOTTED_CURVE_TYPE, color = V_color)
-    ax2.title.set_text(figure_t.gen_title(var + ' - V', start, end))
+    ax_V.set_xlabel(var_disp)
+    ax_V.set_ylabel(r'$V$')
+    ax_V.plot(values, V_ave, DOTTED_CURVE_TYPE, color = V_color)
+    ax_V.set_title(figure_t.gen_title(var_disp + r'$\,-\,V$', start, end))
     
-    return fig1, fig2
+    return ax_U, ax_V
 
 
 
 
-def var_UV2(var1, var2, values1, values2, var_dirs, start = 0.95, end = 1.0, U_color = 'viridis', V_color = 'viridis', rgb_alpha = 1):
+def var_UV2(var1, var2, values1, values2, var_dirs, ax_U = None, ax_V = None, var1_on_xaxis = True, start = 0.95, end = 1.0, U_color = 'viridis', V_color = 'viridis', rgb_alpha = 1):
     '''
     Plot function for test_var2, plot how two variables influence U, V population.
     Make U, V - var1, var2 curves. y-axis is population, x-axis is var2's values, 
@@ -211,6 +219,8 @@ def var_UV2(var1, var2, values1, values2, var_dirs, start = 0.95, end = 1.0, U_c
     - values1:    1D list or np.array, values for var1.
     - values2:    1D list or np.array, values for var2.
     - var_dirs:   return value of test_var2, a 2D list of directories where all data are stored.
+    - ax_U, ax_V:   matplotlib axes to plot on.
+    - var1_on_xaxis:    whether to put var1 on x-axis. Set to False if not.
     - start, end: floats, give an interval of time over which to take average and make plot.
                 For example, start = 0.95, end = 1.0 plots the near-end population (equilibrium, if converged).
     - color:      str, what colors to use for different value of var1. Uses Matplotlib color maps. 
@@ -218,7 +228,7 @@ def var_UV2(var1, var2, values1, values2, var_dirs, start = 0.95, end = 1.0, U_c
     - rgb_alpha:  the alpha value for color. Thr curves might have overlaps, recommend set a small alpha value if so.
     
     Returns:
-    - fig1, fig2: figures for U, V, respectively.
+    - ax_U, ax_V: axes for U, V, respectively.
     '''
 
     # average value of U, V over the interval. One entry for one values of var1, var2
@@ -240,38 +250,59 @@ def var_UV2(var1, var2, values1, values2, var_dirs, start = 0.95, end = 1.0, U_c
 
             U_ave[k1].append(sum(figure_t.ave_interval_1D(simk.U, start_index, end_index)) / NM)
             V_ave[k1].append(sum(figure_t.ave_interval_1D(simk.V, start_index, end_index)) / NM)
-        
-    #### plot ####
+
+    U_ave = np.array(U_ave)
+    V_ave = np.array(V_ave)
     
-    fig1, ax1 = plt.subplots()
-    ax1.set_xlabel(var2)
-    ax1.set_ylabel('U')
-    cmap1 = plt.get_cmap(U_color)
+    #### plot ####
+    if ax_U == None:
+        _, ax_U = plt.subplots()
+    if ax_V == None:
+        _, ax_V = plt.subplots()
 
-    for k1 in range(len(values1)):
-        label = var1 + '=' + str(values1[k1])
-        color_k1 = cmap1(k1 / len(values1))[:3] + (rgb_alpha,)
-        ax1.plot(values2, U_ave[k1], DOTTED_CURVE_TYPE, label = label, color = color_k1)
-    ax1.title.set_text(figure_t.gen_title(var1 + '&' + var2 + ' - U', start, end))
-    ax1.legend(bbox_to_anchor = (1, 1))
+    var1_disp = PATCH_VAR_DISP[var1]
+    var2_disp = PATCH_VAR_DISP[var2]
     
-    fig2, ax2 = plt.subplots()
-    ax2.set_xlabel(var2)
-    ax2.set_ylabel('V')
-    cmap2 = plt.get_cmap(V_color)
 
-    for k1 in range(len(values1)):
-        label = var1 + '=' + str(values1[k1])
-        color_k1 = cmap2(k1 / len(values1))[:3] + (rgb_alpha,)
-        ax2.plot(values2, V_ave[k1], DOTTED_CURVE_TYPE, label = label, color = color_k1)
-    ax2.title.set_text(figure_t.gen_title(var1 + '&' + var2 + ' - V', start, end))
-    ax2.legend(bbox_to_anchor = (1.2, 1))
+    cmap_U = plt.get_cmap(U_color)
+    if var1_on_xaxis:
+        for k in range(len(values2)):
+            label = var2_disp + '=' + str(values2[k])
+            color_k = cmap_U(k / len(values2))[:3] + (rgb_alpha,)
+            ax_U.plot(values1, U_ave[:, k], DOTTED_CURVE_TYPE, label = label, color = color_k)
+            ax_U.set_xlabel(var1_disp)
+    else:
+        for k in range(len(values1)):
+            label = var1_disp + '=' + str(values1[k])
+            color_k = cmap_U(k / len(values1))[:3] + (rgb_alpha,)
+            ax_U.plot(values2, U_ave[k], DOTTED_CURVE_TYPE, label = label, color = color_k)
+            ax_U.set_xlabel(var2_disp)
+    ax_U.set_ylabel(r'$U$')
+    ax_U.set_title(figure_t.gen_title(var1_disp + ',' + var2_disp + r'$\,-\,U$', start, end))
+    ax_U.legend()
     
-    return fig1, fig2
+    cmap_V = plt.get_cmap(V_color)
+    if var1_on_xaxis:
+        for k in range(len(values2)):
+            label = var2_disp + '=' + str(values2[k])
+            color_k = cmap_V(k / len(values2))[:3] + (rgb_alpha,)
+            ax_V.plot(values1, V_ave[:, k], DOTTED_CURVE_TYPE, label = label, color = color_k)
+            ax_V.set_xlabel(var1_disp)
+    else:
+        for k in range(len(values1)):
+            label = var1_disp + '=' + str(values1[k])
+            color_k = cmap_V(k / len(values1))[:3] + (rgb_alpha,)
+            ax_V.plot(values2, V_ave[k], DOTTED_CURVE_TYPE, label = label, color = color_k)
+            ax_V.set_xlabel(var2_disp)
+    ax_V.set_ylabel(r'$V$')
+    ax_V.set_title(figure_t.gen_title(var1_disp + ',' + var2_disp + r'$\,-\,V$', start, end))
+    ax_V.legend()
+    
+    return ax_U, ax_V
 
 
 
-def var_pi1(var, values, var_dirs, start = 0.95, end = 1.0, U_color = 'violet', V_color  = 'yellowgreen'):
+def var_pi1(var, values, var_dirs, ax_U = None, ax_V = None, start = 0.95, end = 1.0, U_color = 'violet', V_color  = 'yellowgreen'):
     '''
     Plot function for test_var1. Plot influence of var on U, V's payoff.
     Make U_pi, V_pi - var curves. y-axis is payoff, x-axis is values of var.
@@ -280,10 +311,11 @@ def var_pi1(var, values, var_dirs, start = 0.95, end = 1.0, U_color = 'violet',
     - var_dirs:   return value of test_var1. A 1D list of directories where all data are stored.
     - var:        str, which variable as tested.
     - values:     1D list or np.array, what values were used.
+    - ax_U, ax_V:   matplotlib axes to plot on.
     - start, end: floats, define an interval of time over which to calculate average payoff and make plots.
     
     Returns:
-    - fig1, fig2: figures for U, V's payoff, respectively.
+    - ax_U, ax_V: axes for U, V, respectively.
     '''
     
     # take average value of payoff over an interval of time
@@ -308,24 +340,28 @@ def var_pi1(var, values, var_dirs, start = 0.95, end = 1.0, U_color = 'violet',
         del simk
             
     #### plot ####
+    if ax_U == None:
+        _, ax_U = plt.subplots()
+    if ax_V == None:
+        _, ax_V = plt.subplots()
+    var_disp = PATCH_VAR_DISP[var]
     
-    fig1, ax1 = plt.subplots()
-    ax1.set_xlabel(var)
-    ax1.set_ylabel('U_pi')
-    ax1.plot(values, U_ave, DOTTED_CURVE_TYPE, color = U_color)
-    ax1.title.set_text(figure_t.gen_title(var + ' - U_pi', start, end))
+    ax_U.set_xlabel(var_disp)
+    ax_U.set_ylabel(r'$\pi_U$')
+    ax_U.plot(values, U_ave, DOTTED_CURVE_TYPE, color = U_color)
+    ax_U.set_title(figure_t.gen_title(var_disp + r'$\,-\,\pi_U$', start, end))
     
-    fig2, ax2 = plt.subplots()
-    ax2.set_xlabel(var)
-    ax2.set_ylabel('V_pi')
-    ax2.plot(values, V_ave, DOTTED_CURVE_TYPE, color = V_color)
-    ax2.title.set_text(figure_t.gen_title(var  + ' - V_pi', start, end))
+    ax_V.set_xlabel(var_disp)
+    ax_V.set_ylabel(r'$\pi_V$')
+    ax_V.plot(values, V_ave, DOTTED_CURVE_TYPE, color = V_color)
+    ax_V.set_title(figure_t.gen_title(var_disp  + r'$\,-\,\pi_V$', start, end))
+
+    return ax_U, ax_V
 
-    return fig1, fig2
 
 
 
-def var_pi2(var1, var2, values1, values2, var_dirs, start = 0.95, end = 1.0, U_color = 'viridis', V_color = 'viridis', rgb_alpha = 1):
+def var_pi2(var1, var2, values1, values2, var_dirs, ax_U = None, ax_V = None, var1_on_xaxis = True, start = 0.95, end = 1.0, U_color = 'viridis', V_color = 'viridis', rgb_alpha = 1):
     '''
     Plot function for test_var2. Plot how var1 and var2 influence payoff.
     Make U_pi, V_pi - var2 curves. y-axis is payoff, x-axis is values of var2,
@@ -335,12 +371,14 @@ def var_pi2(var1, var2, values1, values2, var_dirs, start = 0.95, end = 1.0, U_c
     - var_dirs:           return value of test_var2. A 2D list of directories where all data are stored.
     - var1, var2:         str, what variables were tested.
     - values1, values2:   1D lists or np.array, what values were tested.
+    - ax_U, ax_V:           matplotlib axes to plot on.
+    - var1_on_xaxis:    whether to put var1 on x-axis. Set to False if not.
     - start, end:         floats, define a time inteval over which to make average and make plots.
     - color:              str, Matplotlib color maps.
     - rgb_alpha:          set alpha value for curves.
     
     Returns:
-    - fig1, fig2:         U, V's payoff figures, respectively.
+    - ax_U, ax_V:         U, V's payoff figures, respectively.
     '''
     
     # take average value of payoff over an interval of time
@@ -364,34 +402,54 @@ def var_pi2(var1, var2, values1, values2, var_dirs, start = 0.95, end = 1.0, U_c
             V_ave[k1].append(np.sum(figure_t.ave_interval(simk.V_pi, start_index, end_index)) / NM)
 
             del simk    # manually delete this large object
-        
-    #### plot ####
-    
-    fig1, ax1 = plt.subplots()
-    ax1.set_xlabel(var2)
-    ax1.set_ylabel('U_pi')
-    cmap1 = plt.get_cmap(U_color)
 
-    for k1 in range(len(var_dirs)):
-        label = var1 + '=' + str(values1[k1])
-        color_k1 = cmap1(k1 / len(values1))[:3] + (rgb_alpha,)
-        ax1.plot(values2, U_ave[k1], DOTTED_CURVE_TYPE, label = label, color = color_k1)
-    ax1.title.set_text(figure_t.gen_title(var1 + '&' + var2 + ' - U_pi', start, end))
-    ax1.legend(bbox_to_anchor = (1, 1))
+    U_ave = np.array(U_ave)
+    V_ave = np.array(V_ave)
     
-    fig2, ax2 = plt.subplots()
-    ax2.set_xlabel(var2)
-    ax1.set_ylabel('V_pi')
-    cmap2 = plt.get_cmap(V_color)
+    #### plot ####
+    if ax_U == None:
+        _, ax_U = plt.subplots()
+    if ax_V == None:
+        _, ax_V = plt.subplots()
 
-    for k1 in range(len(var_dirs)):
-        label = var1 + '=' + str(values1[k1])
-        color_k1 = cmap2(k1 / len(values1))[:3] + (rgb_alpha,)
-        ax2.plot(values2, V_ave[k1], DOTTED_CURVE_TYPE, label = label, color = color_k1)
-    ax2.title.set_text(figure_t.gen_title(var1 + '&' + var2 + ' - V_pi', start, end))
-    ax2.legend(bbox_to_anchor = (1.2, 1))
+    var1_disp = PATCH_VAR_DISP[var1]
+    var2_disp = PATCH_VAR_DISP[var2]
+    
+    cmap_U = plt.get_cmap(U_color)
+    if var1_on_xaxis:
+        for k in range(len(values2)):
+            label = var2_disp + '=' + str(values2[k])
+            color_k = cmap_U(k / len(values2))[:3] + (rgb_alpha,)
+            ax_U.plot(values1, U_ave[:, k], DOTTED_CURVE_TYPE, label = label, color = color_k)
+            ax_U.set_xlabel(var1_disp)
+    else:
+        for k in range(len(values1)):
+            label = var1_disp + '=' + str(values1[k])
+            color_k = cmap_U(k / len(values1))[:3] + (rgb_alpha,)
+            ax_U.plot(values2, U_ave[k], DOTTED_CURVE_TYPE, label = label, color = color_k)
+            ax_U.set_xlabel(var2_disp)
+    ax_U.set_ylabel(r'$\pi_{V}$')
+    ax_U.set_title(figure_t.gen_title(var1_disp + ',' + var2_disp + r'$\,-\,\pi_U$', start, end))
+    ax_U.legend()
+    
+    cmap_V = plt.get_cmap(V_color)
+    if var1_on_xaxis:
+        for k in range(len(values2)):
+            label = var2_disp + '=' + str(values2[k])
+            color_k = cmap_V(k / len(values2))[:3] + (rgb_alpha,)
+            ax_V.plot(values1, V_ave[:, k], DOTTED_CURVE_TYPE, label = label, color = color_k)
+            ax_V.set_xlabel(var1_disp)
+    else:
+        for k in range(len(values1)):
+            label = var1_disp + '=' + str(values1[k])
+            color_k = cmap_V(k / len(values1))[:3] + (rgb_alpha,)
+            ax_V.plot(values2, V_ave[k], DOTTED_CURVE_TYPE, label = label, color = color_k)
+            ax_V.set_xlabel(var2_disp)
+    ax_V.set_ylabel(r'$\pi_{U}$')
+    ax_V.set_title(figure_t.gen_title(var1_disp + ',' + var2_disp + r'$\,-\,\pi_V$', start, end))
+    ax_V.legend()
     
-    return fig1, fig2
+    return ax_U, ax_V
 
 
 

piegy/tools/figure_tools.py

@@ -10,6 +10,7 @@ Functions:
 - scale_interval:   scale interval if sim'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
 '''
 
 
@@ -21,12 +22,13 @@ import seaborn as sns
 # default value is [0.125, 0.11, 0.9, 0.88]
 
 
-def heatmap(data, cmap = "Greens", annot = False, fmt = '.3g', title = None, text = None):
+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.
@@ -37,23 +39,26 @@ def heatmap(data, cmap = "Greens", annot = False, fmt = '.3g', title = None, tex
         fig:    Seaborn heatmap.
     '''
 
-    fig, ax = plt.subplots()
+    if ax == None:
+        _, ax = plt.subplots()
+
     if text != None:
-        ax.text(0.63, 0.9, text, size = 10, linespacing = 1.5, transform = plt.gcf().transFigure)
+        ax.text(0.66, 0.9, text, size = 10, linespacing = 1.5, transform = plt.gcf().transFigure)
 
-    ax = sns.heatmap(data, cmap = cmap, annot = annot, fmt = fmt)
-    ax.title.set_text(title)
+    ax = sns.heatmap(data, ax = ax, cmap = cmap, annot = annot, fmt = fmt)
+    ax.set_title(title)
     
-    return fig
+    return ax
 
 
 
-def bar(data, color = "green", xlabel = None, ylabel = None, title = None, text = None):
+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.
@@ -67,44 +72,48 @@ def bar(data, color = "green", xlabel = None, ylabel = None, title = None, text
     N = np.array(data).size
     xaxis = np.array([i for i in range(N)])
     
-    # make figure larger if has more data points
-    fig, ax = plt.subplots()
-    if N > 60:
-        fig.set_size_inches(min(N * 0.12, 9.6), 4.8)
+    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.63, 0.9, text, size = 10, linespacing = 1.5, transform = plt.gcf().transFigure)
+        ax.text(0.66, 0.9, text, size = 10, linespacing = 1.5, transform = plt.gcf().transFigure)
 
     ax.bar(x = xaxis, height = data, color = color)
     ax.set_xlabel(xlabel)
     ax.set_ylabel(ylabel)
-    ax.title.set_text(title)
+    ax.set_title(title)
     
-    return fig
+    return ax
 
 
 
-def scatter(X, Y, color = "orange", alpha = 0.25, xlabel = "x", ylabel = "y", title = None):
+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.
     '''
     
-    fig, ax = plt.subplots()
+    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.title.set_text(title)
+    ax.set_title(title)
 
-    return fig
+    return ax
 
 
 
@@ -121,7 +130,7 @@ def gen_text(ave, std):
     '''
     Generate text about standard deviation info.
     '''
-    text = "ave: " + str(round(ave, 3)) + ", std: " + str(round(std, 3))
+    text = "ave=" + str(round(ave, 2)) + ", std=" + str(round(std, 2))
     return text
 
 
@@ -204,3 +213,26 @@ def scale_interval(interval, compress_itv):
     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/videos.py

@@ -11,22 +11,21 @@ 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.
-- sort_frames:  Put frames in order.
-other not documented here.
+- 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
-import imageio.v2 as imageio
-import re
-from moviepy import VideoFileClip
+from moviepy import ImageSequenceClip
 
 
 # a list of supported figures
@@ -163,15 +162,15 @@ def frame_heatmap_lim(sim, func, frames):
     V_list = []
 
     for i in range(10):
-        U_fig, V_fig = func(sim, start = i / 10, end = (i / 10 + 1 / frames))
+        fig_U, ax_U = plt.subplots()
+        fig_V, ax_V = plt.subplots()
+        ax_U, ax_V = func(sim, ax_U = ax_U, ax_V = ax_V, start = i / 10, end = (i / 10 + 1 / frames))
 
-        U_ax = U_fig.get_axes()[0]
-        U_list.append(U_ax.collections[0].get_clim())
-        V_ax = V_fig.get_axes()[0]
-        V_list.append(V_ax.collections[0].get_clim())
+        U_list.append(ax_U.collections[0].get_clim())
+        V_list.append(ax_V.collections[0].get_clim())
 
-        plt.close(U_fig)
-        plt.close(V_fig)
+        plt.close(fig_U)
+        plt.close(fig_V)
 
     U_clim = get_max_lim(U_list)
     V_clim = get_max_lim(V_list)
@@ -180,49 +179,29 @@ def frame_heatmap_lim(sim, func, frames):
 
 
 
-def sort_frames(images):
-    '''
-    Put frames in order.
-
-    Inputs:
-        images: A list of dirs (frame names)
-    '''
-    numeric_part, non_numeric_part = re.match(r'(\d+) (\D+)', images).groups()
-    return (int(numeric_part), non_numeric_part)
-
-
 
-def make_mp4(dirs, frame_dirs, duration, video_name):
+def make_mp4(video_dir, frame_dir, fps):
     '''
-    Convert frames into a mp4 video.
-
+    Read .png from the frames folder and make into a mp4
     Inputs:
-        dirs:       where to store the video.
-        frame_dirs: where to find frames.
-        duration:   how long the video should be.
-        video_name: name of the video.
+        video_dir:  where to save the video
+        frame_dirs: where to read frames from
+        fps:        frames per second
     '''
 
-    # png to gif
-    images = [img for img in os.listdir(frame_dirs) if img.endswith('.png')] 
-    images.sort(key = sort_frames)
+    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)
 
-    image_list = []
-    for img in images:
-        img_path = os.path.join(frame_dirs, img)
-        image_list.append(imageio.imread(img_path))
-    gif_dirs = dirs + '/temp.gif'
-    imageio.mimsave(gif_dirs, image_list, format = 'gif', duration = duration)
-    
-    # gif to mp4
-    clip = VideoFileClip(gif_dirs)
-    clip.write_videofile(video_name, logger = None)
-    # delete gif
-    os.remove(gif_dirs)
 
 
 
-def make_video(sim, func_name = 'UV_heatmap', frames = 100, speed = 1.25, dpi = 120, U_color = 'Greens', V_color = 'Purples', annot = False, fmt = '.3g', del_frames = False, dirs = 'videos'):
+def make_video(sim, func_name = 'UV_heatmap', frames = 100, 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.
 
@@ -230,8 +209,7 @@ def make_video(sim, func_name = 'UV_heatmap', frames = 100, speed = 1.25, dpi =
     - sim:            a stochastic_model.simulation 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.
-    - speed:          how long every frame should last. Use larger number for slower video.
-    - dpi:            dpi of frames
+    - 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.
@@ -268,7 +246,6 @@ def make_video(sim, func_name = 'UV_heatmap', frames = 100, speed = 1.25, dpi =
     if os.path.exists(V_frame_dirs):
         file_t.del_dirs(V_frame_dirs)
     os.makedirs(V_frame_dirs)
-
         
     #### for loop ####
     
@@ -277,30 +254,31 @@ def make_video(sim, func_name = 'UV_heatmap', frames = 100, speed = 1.25, dpi =
             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:
-            U_fig, V_fig = func(sim, U_color, V_color, start = i / frames, end = (i + 1) / frames, annot = annot, fmt = fmt)
+            ax_U, ax_V = func(sim, 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:
-            U_fig, V_fig = func(sim, U_color, V_color, start = i / frames, end = (i + 1) / frames)
-        U_ax = U_fig.get_axes()[0]
-        V_ax = V_fig.get_axes()[0]
+            ax_U, ax_V = func(sim, 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:
-            U_ax.collections[0].set_clim(U_clim)
-            V_ax.collections[0].set_clim(V_clim)
+            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
-            U_ax.set_ylim(U_ylim)
-            V_ax.set_ylim(V_ylim)
+            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
-                U_ax.set_xlim(U_xlim)
-                V_ax.set_xlim(V_xlim)
+                ax_V.set_xlim(U_xlim)
+                ax_V.set_xlim(V_xlim)
 
-        U_fig.savefig(U_frame_dirs + '/' + str(i) + ' U' + '.png', dpi = dpi)
-        V_fig.savefig(V_frame_dirs + '/' + str(i) + ' V' + '.png', dpi = dpi)
+        fig_U.savefig(U_frame_dirs + '/' + 'U_frame_' + str(i) + '.png', pad_inches = 0.25)
+        fig_V.savefig(V_frame_dirs + '/' + 'V_frame_' + str(i) + '.png', pad_inches = 0.25)
         
-        plt.close(U_fig)
-        plt.close(V_fig)
+        plt.close(fig_U)
+        plt.close(fig_V)
         
     #### for loop ends ####
     
@@ -308,8 +286,8 @@ def make_video(sim, func_name = 'UV_heatmap', frames = 100, speed = 1.25, dpi =
     print('making mp4...      ', end = '\r')
     
     # make videos based on frames
-    make_mp4(dirs, U_frame_dirs, frames * speed, dirs + '/U-' + func_name + '.mp4')
-    make_mp4(dirs, V_frame_dirs, frames * speed, dirs + '/V-' + func_name + '.mp4')
+    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)

{piegy-1.0.0.dist-info → piegy-1.1.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: piegy
-Version: 1.0.0
+Version: 1.1.1
 Summary: Payoff-Driven Stochastic Spatial Model for Evolutionary Game Theory
 Author-email: Chenning Xu <cxu7@caltech.edu>
 License: BSD 3-Clause License
@@ -49,7 +49,6 @@ Description-Content-Type: text/markdown
 License-File: LICENSE.txt
 Requires-Dist: numpy
 Requires-Dist: matplotlib
-Requires-Dist: imageio>=2.37.0
 Requires-Dist: moviepy>=2.1.1
 Requires-Dist: seaborn>=0.13.2
 Dynamic: license-file
@@ -62,20 +61,22 @@ Provides a stochastic spatial model for simulating the interaction and evolution
 
 ## Installation
 
+To install *piegy*, run the following in terminal:
+
 ```bash
 pip install piegy
 ```
 
 ## Documentation and Source
 
-See source code at: [GitHub-piegy repo](https://github.com/Chenning04/piegy.git). 
-The *piegy* documentation at: [piegy docs](https://piegy.readthedocs.io/en/). 
+See source code at: [piegy GitHub repository](https://github.com/Chenning04/piegy.git). 
+The *piegy* documentation at: [piegy Documentation](https://piegy.readthedocs.io/en/). 
 
 ## How the Model Works
 
-Our model can be summarized as "classical game theory endowed with a spatial structure and payoff-driven migration rules". Consider two species, predators and preys (denote by *U* and *V*), in a rectangular region. We divide the region into N by M patches and simulate their interaction within a patch by classical game theory (i.e., payoff matrices and carrying capacity). Interactions across patches are simulated by payoff-driven migration rules. An individual migrates to a neighboring patch with probability weighted by payoff in the neighbors.
+Our model can be summarized as "classical evolutionary game theory endowed with spatial structure and payoff-driven migration rules". Consider two species, predators and preys (denoted by *U* and *V*), in a rectangular region. We divide the region into N by M patches and simulate their interaction within a patch by classical game theory (i.e., payoff matrices and carrying capacity). Interactions across patches are simulated by payoff-driven migration rules. An individual migrates to a neighboring patch with probability weighted by payoff in the neighbors.
 
-We use the Gillepie algorithm as the fundamental event-selection algorithm. At each time step, one event is selected and let happen; and step sizes are continuous, dependent on the current state in the space. Data are recorded every some specified time interval.
+We use the Gillepie algorithm as the fundamental event-selection algorithm. At each time step, one event is selected and let happen; and the step size is continuous, dependent on the current state in the space. Data are recorded every some specified time interval.
 
 ## Analytic Tools
 
@@ -95,10 +96,11 @@ dynamics = figures.UV_dyna(sim)
 U_hmap, V_hmap = figures.UV_heatmap(sim)
 ```
 
-The figures reveal the population dynamics and steady state distribution.
+The figures reveal the population dynamics and steady state population distribution.
 
 
 ## Acknowledgments
 
 - Thanks Professor Daniel Cooney at University of Illinois Urbana-Champaign. This package is developed alongside a project with Prof. Cooney and received enormous help from him.
 - Special thanks to the open-source community for making this package possible.
+

piegy-1.1.1.dist-info/RECORD

@@ -0,0 +1,16 @@
+piegy/__init__.py,sha256=Lrh6NegSvo6LOCXg_tBTu804eicnHpQY2zmE0FchjKE,3241
+piegy/__version__.py,sha256=8IfclnXlmVE5p-odcoFI_9v7RIOmz2Cxd8Dvx1dMjNc,526
+piegy/analysis.py,sha256=1cF06igQMGJGVjLiyhtgepGk8fYhzvL0orI48tOK1qY,8713
+piegy/data_tools.py,sha256=lNFOX4H0o_WfRSNQoC9QGCK66-IdrZ0zCnz12N-Ael4,3599
+piegy/figures.py,sha256=xR4NKuHc-7E3ld6NjyAAVcJG8CtIoKoM4mqd8qPOXsg,19063
+piegy/model.py,sha256=NkP_Apnun3MhATSihR4_irnA-8p4FW879xQjCXVEPMA,45502
+piegy/test_var.py,sha256=PDtCCou-kFXumEe2EdAQNKJohwHH58ifLSsp5xQraVI,23532
+piegy/videos.py,sha256=vPl1aTIOaWxIUsxX5yL3JPiuf-Um0bPo7qYIHZP0bo0,10143
+piegy/tools/__init__.py,sha256=eYOl_HJHDonYexfrmKh3koOlxvtSo46vH6jHvCEEB4k,300
+piegy/tools/figure_tools.py,sha256=tAI5sBMsxN_1BN5_G_3UoMcDeYTj2DXr8l4HpPl9FSs,6961
+piegy/tools/file_tools.py,sha256=ncxFWeHfIE-GYLQlOrahFlhBgqPyuY3R5_93fpQeCEs,630
+piegy-1.1.1.dist-info/licenses/LICENSE.txt,sha256=wfzEht_CxOcfGGmg3f3at4mWJb9rTBjA51mXLl_3O3g,1498
+piegy-1.1.1.dist-info/METADATA,sha256=qawzMEDrKZITzujlqMf_9WJoEklAGFYcOw3I2u8plbQ,5097
+piegy-1.1.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+piegy-1.1.1.dist-info/top_level.txt,sha256=k4QLYL8PqdqDuy95-4NZD_FVLqJDsmq67tpKkBn4vMw,6
+piegy-1.1.1.dist-info/RECORD,,

{piegy-1.0.0.dist-info → piegy-1.1.1.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (80.4.0)
+Generator: setuptools (80.9.0)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

piegy-1.0.0.dist-info/RECORD

@@ -1,16 +0,0 @@
-piegy/__init__.py,sha256=Lrh6NegSvo6LOCXg_tBTu804eicnHpQY2zmE0FchjKE,3241
-piegy/__version__.py,sha256=H58EUm6wBWYw2Q4gXk5PznfRrEry4c4zk7c--daUx_8,365
-piegy/analysis.py,sha256=1cF06igQMGJGVjLiyhtgepGk8fYhzvL0orI48tOK1qY,8713
-piegy/data_tools.py,sha256=dXESWglAXN4q79ZG8wtpZftJoiiaOJFU7NSaip1ZTq0,3518
-piegy/figures.py,sha256=46Vg6AsWswF-juGbyndaRaChrdVuzPbpRQQd9grTWfk,17842
-piegy/model.py,sha256=2Y49NDzdtog9effo9Lr_-psu5NMBF3iH9BeE8zhCHZ8,45700
-piegy/test_var.py,sha256=fk_e0Hko6zNfsOx-SMnnEcsgeu6fx0rQcdfuXGceY7U,20892
-piegy/videos.py,sha256=vW0F9WYDvPVJM0fHM_0Vi9Y8ZSVfJBcuCLJ2lLXeZ2U,10410
-piegy/tools/__init__.py,sha256=eYOl_HJHDonYexfrmKh3koOlxvtSo46vH6jHvCEEB4k,300
-piegy/tools/figure_tools.py,sha256=qhuMPZn2krEzjhjmAd6ac9QVBeCWUkSdCGOz7RASAhA,5931
-piegy/tools/file_tools.py,sha256=ncxFWeHfIE-GYLQlOrahFlhBgqPyuY3R5_93fpQeCEs,630
-piegy-1.0.0.dist-info/licenses/LICENSE.txt,sha256=wfzEht_CxOcfGGmg3f3at4mWJb9rTBjA51mXLl_3O3g,1498
-piegy-1.0.0.dist-info/METADATA,sha256=hn5JGkA5jBMN2toZ0xPdJ-0munvYBE-BzcOMKWb-d-g,5035
-piegy-1.0.0.dist-info/WHEEL,sha256=DnLRTWE75wApRYVsjgc6wsVswC54sMSJhAEd4xhDpBk,91
-piegy-1.0.0.dist-info/top_level.txt,sha256=k4QLYL8PqdqDuy95-4NZD_FVLqJDsmq67tpKkBn4vMw,6
-piegy-1.0.0.dist-info/RECORD,,