Modified the argument name to sort_by_priority

primo/data_parser/well_data.py

@@ -298,7 +298,7 @@ def get_fully_partitioned_data(self):
 
         return None
 
-    def get_high_priority_wells(self, num_wells: int, sort: bool = True):
+    def get_high_priority_wells(self, num_wells: int, sort_by_priority: bool = True):
         """
         Returns the top n wells by priority
 
@@ -307,7 +307,7 @@ def get_high_priority_wells(self, num_wells: int, sort: bool = True):
         num_wells : int
             Returns the top `num_wells` number of wells
 
-        sort : bool, default = True
+        sort_by_priority : bool, default = True
             If True, arranges the wells in descending order of priority.
             If False, the wells are arranged as per their index.
         """
@@ -321,7 +321,7 @@ def get_high_priority_wells(self, num_wells: int, sort: bool = True):
             self.data.sort_values(cn, ascending=False).head(num_wells).index.to_list()
         )
 
-        if sort:
+        if sort_by_priority:
             # Constructed sub-data is ordered as per index, not priority.
             # So, sorting the data again
             hp_wells.data = hp_wells.data.sort_values(cn, ascending=False)

primo/utils/clustering_utils.py

@@ -51,7 +51,7 @@ def distance_matrix(
 
     Returns
     -------
-    np.ndarray
+    pd.DataFrame
         Distance matrix to be used for the agglomerative
         clustering method
 
@@ -110,6 +110,9 @@ def perform_clustering(wd: WellData, distance_threshold: float = 10.0):
 
     Parameters
     ----------
+    wd : WellData
+        Object containing the information on all wells
+
     distance_threshold : float, default = 10.0
         Threshold distance for breaking clusters
 
@@ -152,7 +155,8 @@ def perform_clustering(wd: WellData, distance_threshold: float = 10.0):
 
 def get_pairwise_metrics(wd: WellData, list_wells: list) -> pd.DataFrame:
     """
-    Returns pairwise metric values for all well pairs in each clusters
+    Returns pairwise metric values for all possible pairs of wells in
+    `list_wells`.
 
     Parameters
     ----------