piegy 1.0.0__py3-none-any.whl

piegy/test_var.py

@@ -0,0 +1,525 @@
+'''
+Test how certain variables influence simulation result.
+
+Public Functions:
+
+Test influence of one variable:
+- test_var1:    Test how a certain patch variable (mu, w, kappa) influences results.
+                Run simulations across different values of that variable and save data.
+- var_UV1:      Plot how U, V change across different values of a specified variable.
+- var_pi1:      Plot how U_pi, V_pi change across different values of a specified variable.
+                    
+
+Test influence of two variables:
+- test_var2:    Test how two patch variables influence results.
+                Run simulations across different values of the two specified variables and save data.
+- var_UV2:      Plot how U, V change across the two variables.
+                x-axis is values of the variable 2, y-axis is U or V. Values of variable 1 are shown by multiple curves.
+- var_pi2:      Plot how U_pi & V_pi change across the two variables.
+
+
+Additional tools (can be called directly):
+- get_dirs1:   test_var1 returns directories in a certain format based on var and values. This function mimics that format and 
+                        returns a 1D list of directories, where all data are stored.
+                        This 1D list is intended to pass to var_UV1 and var_pi1 directly.
+    
+- var_convergence1: Find the simulatoin results of test_var1 that diverge.
+
+- get_dirs2:   test_var2 returns directories in a certain format based on vars and values. This function mimics that format and
+                    returns a 2D list of directories, where all data are stored.
+                    This 2D list is intended to pass to var_UV2 and var_pi2 directly.
+- var_convergence2: Find the simulatoin results of test_var2 that diverge.
+'''
+
+from . import model as model
+from .tools import figure_tools as figure_t
+from . import analysis as analysis
+from . import data_tools as data_t
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+
+
+# curve type used by var_UV, var2_UV, var_pi, var2_pi
+# can be 'o-', 'x-', ...
+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}
+
+
+
+def test_var1(sim, var, values, dirs, compress_itv = None, predict_runtime = False):
+    '''
+    Test the influence of one patch variable on simulation results.
+
+    Inputs::
+    - sim:                a simulation object. All tests will use parameters of sim, except for the variable to test.
+    - var:                str, which patch variable to test. e.g. var can be 'mu1', 'w2', 'kappa2', ...
+    - values:             1D np.array or list, what values of var to test.
+    - dirs:               str, where to save data.
+    - compress_itv:         int, whether to reduce data size (if not 1), passed to model.simulation.compress_data function.
+    - predict_runtime:    bool, whether to predict how much time left for each test.          
+
+    Returns::
+    - var_dirs:           a 1D list of directories (str) where all the data are stored. 
+                        Directories have form 'mu1=0.1'. 
+                        NOTE: if you accidentally lost this return value, you can retrieve it by get_dirs1.
+    '''
+
+    if var not in PATCH_VAR_DICT.keys():
+        raise ValueError('Please enter a valid patch variable: mu1, mu2, w1, w2, kappa1, or kappa2')
+
+    # var_dirs is where data will be stored
+    var_dirs = []
+    
+    for k in range(len(values)):
+        sim2 = sim.copy(copy_data = False)
+        current_var_str = var + '=' + str(values[k])    # e.g., 'mu1=0.1'
+        var_dirs.append(dirs + '/' + current_var_str)
+        
+        for i in range(sim.N):
+            for j in range(sim.M):
+                sim2.P[i][j][PATCH_VAR_DICT[var]] = values[k]
+
+        try:
+            model.run(sim2, predict_runtime, message = current_var_str + ', ')
+            if compress_itv != None:
+                sim2.compress_data(compress_itv)
+            data_t.save_data(sim2, var_dirs[k], print_msg = False)
+        except OverflowError:
+            print(current_var_str + ' raised OverflowError, skipped')
+            continue
+
+    return var_dirs
+
+
+
+def test_var2(sim, var1, var2, values1, values2, dirs, compress_itv = None, predict_runtime = False):
+    '''
+    Two-variable version of test_var1. Test the influence of two varibles on simulation results.
+
+    Inputs:
+    - sim:      a simulation object. All tests will use the parameters of sim, except for the two vars to be tested.
+    - var1:     str, the first variable to test.
+    - var2:     str, the second variable to test.
+    - values1:  1D list or np.array, values for var1.
+    - values2:  1D list or np.array, values for var2.
+    - dirs, compress_itv, scale_maxtime, predict_runtime: same as in test_var1
+
+    Returns:
+    - var_dirs:   2D list of directories, where all the data are stored.
+                Directories have form 'mu1=0.1, mu2=0.2'.
+                NOTE: if you accidentally lost this return value, you can retrieve by get_dirs2.
+    '''
+
+    if (var1 not in PATCH_VAR_DICT.keys()) or (var2 not in PATCH_VAR_DICT.keys()):
+        raise ValueError('Please enter a valid patch variable: mu1, mu2, w1, w2, kappa1, or kappa2')
+
+    var_dirs = [[] for k1 in range(len(values1))]
+
+    for k1 in range(len(values1)):
+        for k2 in range(len(values2)):
+            sim2 = sim.copy(copy_data = False)
+            current_var_str = var1 + '=' + str(values1[k1]) + ', ' + var2 + '=' + str(values2[k2])   # e.g., mu1=0.1, mu2=0.2
+            var_dirs[k1].append(dirs + '/' + current_var_str)
+
+            for i in range(sim.N):
+                for j in range(sim.M):
+                    sim2.P[i][j][PATCH_VAR_DICT[var1]] = values1[k1]
+                    sim2.P[i][j][PATCH_VAR_DICT[var2]] = values2[k2]
+
+            try:
+                model.run(sim2, predict_runtime, message = current_var_str + ', ')
+                if compress_itv != None:
+                    sim2.compress_data(compress_itv)
+                data_t.save_data(sim2, var_dirs[k1][k2], print_msg = False)
+            except OverflowError:
+                print(current_var_str + ' raise OverflowError, skipped')
+                continue
+
+    return var_dirs
+
+
+
+
+def var_UV1(var, values, var_dirs, 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.
+
+    Inputs:
+    - 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.
+    - 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.
+    '''
+
+    # average value of U, V over the interval. One entry for one value of var
+    U_ave = []
+    V_ave = []
+
+    for k in range(len(var_dirs)):
+        try:
+            simk = data_t.read_data(var_dirs[k])
+        except FileNotFoundError:
+            print(var + '=' + str(values[k]) + ' not found, skipped')
+            U_ave.append(None)
+            V_ave.append(None)
+            continue
+        start_index = int(simk.max_record * start)
+        end_index = int(simk.max_record * end)
+        NM = int(simk.N * simk.M)
+
+        U_ave.append(sum(figure_t.ave_interval_1D(simk.U, start_index, end_index)) / NM)
+        V_ave.append(sum(figure_t.ave_interval_1D(simk.V, start_index, end_index)) / NM)
+        
+    #### plot ####
+    
+    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))
+    
+    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))
+    
+    return fig1, fig2
+
+
+
+
+def var_UV2(var1, var2, values1, values2, var_dirs, 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, 
+    and var1's values are represented by different curves, one curve corresponds to one value of var1.
+
+    Inputs:
+    - var1:       str, the first variable tested.
+    - var2:       str, the second variable tested.
+    - 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.
+    - 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. 
+                See available color maps at: https://matplotlib.org/stable/users/explain/colors/colormaps.html
+    - 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.
+    '''
+
+    # average value of U, V over the interval. One entry for one values of var1, var2
+    U_ave = [[] for k1 in range(len(var_dirs))]
+    V_ave = [[] for k1 in range(len(var_dirs))]
+
+    for k1 in range(len(var_dirs)):
+        for k2 in range(len(var_dirs[k1])):
+            try:
+                simk = data_t.read_data(var_dirs[k1][k2])
+            except FileNotFoundError:
+                print(var1 + '=' + str(values1[k1]) + ', ' + var2 + '=' + str(values2[k2]) + ' not found, skipped')
+                U_ave[k1].append(None)
+                V_ave[k1].append(None)
+                continue
+            start_index = int(simk.max_record * start)
+            end_index = int(simk.max_record * end)
+            NM = int(simk.N * simk.M)
+
+            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 ####
+    
+    fig1, ax1 = plt.subplots()
+    ax1.set_xlabel(var2)
+    ax1.set_ylabel('U')
+    cmap1 = plt.get_cmap(U_color)
+
+    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))
+    
+    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))
+    
+    return fig1, fig2
+
+
+
+def var_pi1(var, values, var_dirs, 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.
+
+    Inputs:
+    - 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.
+    - 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.
+    '''
+    
+    # take average value of payoff over an interval of time
+    U_ave = []
+    V_ave = []
+    
+    for k in range(len(var_dirs)):
+        try:
+            simk = data_t.read_data(var_dirs[k])
+        except FileNotFoundError:
+            print(var + '=' + str(values[k]) + ' not found, skipped')
+            U_ave.append(None)
+            V_ave.append(None)
+            continue
+        start_index = int(simk.max_record * start)
+        end_index = int(simk.max_record * end)
+        NM = int(simk.N * simk.M)
+
+        U_ave.append(np.sum(figure_t.ave_interval(simk.U_pi, start_index, end_index)) / NM)
+        V_ave.append(np.sum(figure_t.ave_interval(simk.V_pi, start_index, end_index)) / NM)
+
+        del simk
+            
+    #### plot ####
+    
+    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))
+    
+    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))
+
+    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):
+    '''
+    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,
+    var1 is represented by different curves. One curve corresponds to one value of var1.
+
+    Inputs:
+    - 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.
+    - 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.
+    '''
+    
+    # take average value of payoff over an interval of time
+    U_ave = [[] for k1 in range(len(var_dirs))]
+    V_ave = [[] for k1 in range(len(var_dirs))]
+
+    for k1 in range(len(var_dirs)):
+        for k2 in range(len(var_dirs[k1])):
+            try:
+                simk = data_t.read_data(var_dirs[k1][k2])
+            except FileNotFoundError:
+                print(var1 + '=' + str(values1[k1]) + ', ' + var2 + '=' + str(values2[k2]) + ' not found, skipped')
+                U_ave[k1].append(None)
+                V_ave[k1].append(None)
+                continue
+            start_index = int(simk.max_record * start)
+            end_index = int(simk.max_record * end)
+            NM = int(simk.N * simk.M)
+
+            U_ave[k1].append(np.sum(figure_t.ave_interval(simk.U_pi, start_index, end_index)) / NM)
+            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))
+    
+    fig2, ax2 = plt.subplots()
+    ax2.set_xlabel(var2)
+    ax1.set_ylabel('V_pi')
+    cmap2 = plt.get_cmap(V_color)
+
+    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))
+    
+    return fig1, fig2
+
+
+
+
+def get_dirs1(var, values, dirs):
+    '''
+    Mimics the return value format of test_var1 and returns a 1D list of directories where test_var1 saved data.
+    Intended usage: retrieve the directories if you accidentally lost the return value of test_var1.
+
+    Inputs:
+    - var:        what variable was tested.
+    - values:     what values were testsed.
+    - dirs:  the directory you passed to test_var1 as 'dirs' parameter. 
+                Essentially the root directories where data were stored.
+
+    Returns:
+    - var_dirs:   a 1D list of directories where data were stored. Has the same format as the return value of test_var1.
+    '''
+
+    var_dirs = []
+    values_sorted = sorted(values)
+    if dirs[-1] != '/':
+        dirs += '/'
+
+    for val in values_sorted:
+            var_dirs.append(dirs + var + '=' + str(val))
+
+    return var_dirs
+
+
+
+
+def var_convergence1(var_dirs, interval = 20, start = 0.8, fluc = 0.07):
+    '''
+    Find the simulatoin results of test_var1 that diverge. 
+
+    Inputs:
+    - var_dirs:       Return value of test_var1
+    - interval:       int. One of the inputs of model_analysis.check_convergence.
+                    The size of interval to take average over.
+    - start:          (0,1) float. One of the inputs of model_analysis.check_convergence. 
+                    Convergence is expected to start from at least this point.
+    - fluc:           (0,1) float. One of the inputs of model_analysis.check_convergence. 
+                    Expect the difference between any two small intervals (a quotient-form difference) should be less than fluc.
+
+    Returns:
+    - diverge_list:   A list directories where the simulation data diverge.
+    '''
+
+    diverge_list = []
+
+    for dirs in var_dirs:
+        try:
+            sim = data_t.read_data(dirs)
+        except FileNotFoundError:
+            print(dirs + ' data not found, skipped')
+            continue
+        if not analysis.check_convergence(sim, interval, start, fluc):
+            diverge_list.append(dirs)
+
+    return diverge_list
+
+
+
+
+def get_dirs2(var1, var2, values1, values2, dirs):
+    '''
+    Mimics the return value format of test_var2 and returns a 2D list of directories where test_var2 saved data.
+    Intended usage: retrieve the directories if you accidentally lost the return value of test_var2.
+
+    Inputs:
+    - var1, var2:         what variables were tested.
+    - values1, values2:   what values were testsed.
+    - dirs:          the directory you passed to test_var2 as 'dirs' parameter. 
+                        Essentially the root directories where data were stored.
+
+    Returns:
+    - var_dirs:   a 2D list of directories where data were stored. Has the same format as the return value of test_var2.
+    '''
+
+    var_dirs = [[] for i in range(len(values1))]
+    values1_sorted = sorted(values1)
+    values2_sorted = sorted(values2)
+    if dirs[-1] != '/':
+        dirs += '/'
+
+    for i in range(len(values1)):
+        for j in range(len(values2)):
+            v1 = values1_sorted[i]
+            v2 = values2_sorted[j]
+            dirs_ij = dirs + var1 + '=' + str(v1) + ', ' + var2 + '=' + str(v2)
+            var_dirs[i].append(dirs_ij)
+
+    return var_dirs
+
+
+
+
+def var_convergence2(var_dirs, interval = 20, start = 0.8, fluc = 0.07):
+    '''
+    Find the simulatoin results of test_var2 that diverge.
+
+    Inputs:
+    - var_dirs:       Return value of test_var2
+    - interval:       int. One of the inputs of model_analysis.check_convergence.
+                    The size of interval to take average over.
+    - start:          (0,1) float. One of the inputs of model_analysis.check_convergence. 
+                    Convergence is expected to start from at least this point.
+    - fluc:           (0,1) float. One of the inputs of model_analysis.check_convergence. 
+                    Expect the difference between any two small intervals (a quotient-form difference) should be less than fluc.
+
+    Returns:
+    - diverge_list:   A list directories where the simulation data diverge.
+    '''
+
+    diverge_list = []
+
+    for sublist in var_dirs:
+        for dirs in sublist:
+            try:
+                sim = data_t.read_data(dirs)
+            except FileNotFoundError:
+                print(dirs + ' data not found, skipped')
+                continue
+            if not analysis.check_convergence(sim, interval, start, fluc):
+                diverge_list.append(dirs)
+
+    return diverge_list
+
+
+

piegy/tools/__init__.py

@@ -0,0 +1,15 @@
+'''
+Init file for .tools sub-package.
+
+Contains three tool kits:
+    - data_tools:   intended for both direct usage and private helper module.
+    - file_tools:   private use.
+    - figure_tools: private use.
+
+    
+Last updated: Mar 27, 2025
+'''
+
+from . import file_tools
+from . import figure_tools
+