Added docstrings to generation modules

Docstrings have been added to the methods in the group generation, individual status generation, and test result generation modules.

group_testing/generate_groups.py

@@ -26,22 +26,29 @@
 # import plotting and data-manipulation tools
 # import matplotlib.pyplot as plt
 
-"""
-Unit testing to be added later
-"""
-# class TestMeasurementMatrix(unittest.TestCase):
-#
-#    def test_make_graph(self):
-#        gen_measurement_matrix(opts)
-
 
-"""
-Function to generate and return the matrix
-"""
 
 
 def gen_measurement_matrix(seed=0, N=1000, m=100, group_size=4, max_tests_per_individual=4, verbose=False,
                            graph_gen_method='no_multiple', plotting=False, saving=True, run_ID='debugging'):
+    """
+    Function to generate and return the group testing matrix based on near-doubly-regular design 
+
+    Parameters:
+        seed (int): Seed for random number generation
+        N (int): Population size
+        m (int): Number of group tests
+        group_size (int): Size of the groups
+        max_tests_per_individual (int): Maximum number of tests allowed per individual
+        verbose (bool): Flag for turning on debugging print statements
+        graph_gen_method (str): Method for igraph to use in generating the graph
+        plotting (bool): Flag for turning on plotting 
+        saving (bool): Flag for turning on saving the result
+        run_ID (str): String for specifying name of run
+
+    Returns:
+        A (binary numpy array): The group testing matrix
+    """
 
     # set the seed used for graph generation to the options seed
     random.seed(seed)
@@ -203,8 +210,10 @@ def gen_measurement_matrix(seed=0, N=1000, m=100, group_size=4, max_tests_per_in
     return A
 
 
-# main method for testing
 if __name__ == '__main__':
+    """
+    Main method for testing
+    """
 
     # print igraph version
     print("Loaded igraph version {}".format(igraph.__version__))

group_testing/generate_individual_status.py

@@ -15,10 +15,19 @@
     formatter=dict(float=lambda x: "%.3g" % x))
 import scipy.io as sio
 
-"""
-Function to generate the infected status of individuals (a vector)
-"""
 def gen_status_vector(seed=0, N=1000, s=10, verbose=False):
+    """
+    Function to generate the infected status of individuals (a vector)
+
+    Parameters:
+        seed (int): Seed for random number generation
+        N (int): Population size
+        s (int): Number of infected individuals
+        verbose (bool): Flag for turning on debugging print statements
+
+    Returns:
+        u (binary numpy array): The status vector
+    """
 
     # set the seed used for status generation
     local_random = random.Random()
@@ -42,6 +51,9 @@ def gen_status_vector(seed=0, N=1000, s=10, verbose=False):
     return u
 
 if __name__ == '__main__':
+    """
+    Main method for testing
+    """
 
     # options for plotting, verbose output, saving, seed
     opts = {}

group_testing/generate_test_results.py

@@ -20,14 +20,28 @@
                     formatter=dict(float=lambda x: "%.3g" % x))
 import scipy.io as sio
 
-"""
-Function to generate the results of the tests from the measurement matrix A 
-and the status vector u
-"""
-
-
 def gen_test_vector(A, u, seed=0, m=100, verbose=False, test_noise_methods=[], binary_symmetric_noise_prob=0.26,
                     theta_l=0, theta_u=0.0625, permutation_noise_prob=0.01):
+    """
+    Function to generate the results of the tests from the measurement matrix A 
+    and the status vector u in both the noiseless and noisy setting
+
+    Parameters:
+        A (binary numpy array): The group testing matrix
+        u (binary numpy array): The individual status vector
+        seed (int): Seed for random number generation
+        m (int): Number of group tests
+        verbose (bool): Flag for turning on debugging print statements
+        test_noise_methods (str list): The list of noise models to apply
+        binary_symmetric_noise_prob (float): The noise level for the binary symmetric noise model
+        theta_l (float): The lower bound for the threshold noise model
+        theta_u (float): The upper bound for the threshold noise model
+        permutation_noise_prob (float): The noise level for the permutation noise model
+
+    Returns:
+        b (binary numpy array): The vector of results of the group tests
+    """
+
     # set the seed used for test result noise
     random.seed(seed)
     np_random = np.random.RandomState(seed)
@@ -189,6 +203,9 @@ def gen_test_vector(A, u, seed=0, m=100, verbose=False, test_noise_methods=[], b
 
 
 if __name__ == '__main__':
+    """
+    Main method for testing
+    """
 
     # options for plotting, verbose output, saving, seed
     opts = {}