From beca3b56182434299d053f9ed8ea1343027198d1 Mon Sep 17 00:00:00 2001
From: Nick Dexter <dexter.nicholas@gmail.com>
Date: Thu, 29 Apr 2021 16:57:18 -0700
Subject: [PATCH] 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            | 33 +++++++++++++--------
 group_testing/generate_individual_status.py | 18 +++++++++--
 group_testing/generate_test_results.py      | 29 ++++++++++++++----
 3 files changed, 59 insertions(+), 21 deletions(-)

diff --git a/group_testing/generate_groups.py b/group_testing/generate_groups.py
index 85ea1b3..4bef558 100644
--- a/group_testing/generate_groups.py
+++ b/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__))
diff --git a/group_testing/generate_individual_status.py b/group_testing/generate_individual_status.py
index 4f8827d..4d6a305 100644
--- a/group_testing/generate_individual_status.py
+++ b/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 = {}
diff --git a/group_testing/generate_test_results.py b/group_testing/generate_test_results.py
index 852ea68..10b5435 100644
--- a/group_testing/generate_test_results.py
+++ b/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 = {}