Refactored map_utils file (#61)

* Refactored map_utils file

* Updated get_cluster_colors in map_utils

* Updated map utils class based on feedback

* Updated tests for map_utils
Made addition of shapefile optional to map_utils

* Added test for visualize_campaign

* Removed pylint disable from 2 files. Additional minor changes made from review comments.

* Removed pylint disable from result_parser file

* Added test for get_high_priority_wells in well data tests
Updated functions usage in demo
Updated one test in map_utils

* Fixed bug in adding shapefile to map object

---------

Co-authored-by: Yash Puranik <89143090+apd-ypuranik@users.noreply.github.com>

.github/wordlist.txt

@@ -31,6 +31,7 @@ Submetric
 Submetrics
 Uncomment
 USD
+ValueError
 Vbox
 acs
 arg
@@ -58,6 +59,7 @@ dtype
 dunder
 env
 fixme
+geopandas
 geopoint
 geospatial
 glpk
@@ -66,6 +68,7 @@ inplace
 ipynb
 ipywidgets
 iterable
+itertuples
 json
 libs
 monkeypatch

.stagedfright/checks.py

@@ -32,7 +32,7 @@ def test_allowfile_matches_if_present(staged: StagedFile, allowfile: AllowFile):
     "primo/data_parser/metric_data.py": 15,
     "primo/data_parser/well_data.py": 40,
     "primo/data_parser/tests/test_metric_data.py": 50,
-    "primo/data_parser/tests/test_well_data.py": 237,
+    "primo/data_parser/tests/test_well_data.py": 241,
     "primo/data_parser/tests/test_well_data_columns.py": 13,
     "primo/data_parser/tests/test_well_data_common_methods.py": 73,
     "docs/conf.py": 4,

primo/data_parser/tests/test_well_data.py

@@ -463,6 +463,22 @@ def test_both_production_type(
     assert len(gas_oil_wells) == 2
 
 
+def test_get_high_priority_wells(get_column_names):
+    """Test the get_high_priority_wells method"""
+    filename, col_names = get_column_names
+    col_names.priority_score = "Priority Score [0-100]"
+    df = pd.read_csv(filename)
+    df[col_names.priority_score] = np.linspace(100, 0, len(df))
+    wd = WellData(data=df, column_names=col_names)
+    top_wells = wd.get_high_priority_wells(5)
+
+    assert top_wells is not None
+    assert top_wells.data.shape[0] == 5
+    assert top_wells.data[col_names.priority_score].tolist() == sorted(
+        top_wells.data[col_names.priority_score], reverse=True
+    )
+
+
 def test_well_partitioning(
     tmp_path, get_column_names
 ):  # pylint: disable=too-many-statements

primo/data_parser/well_data.py

@@ -272,6 +272,20 @@ def get_fully_partitioned_data(self):
 
         return None
 
+    def get_high_priority_wells(self, num_wells: int):
+        """Returns the top n wells by priority"""
+
+        if not hasattr(self._col_names, "priority_score"):
+            LOGGER.warning("Returning None, since priority scores are not available!")
+            return None
+
+        well_list = (
+            self.data.sort_values("Priority Score [0-100]", ascending=False)
+            .head(num_wells)
+            .index.to_list()
+        )
+        return self._construct_sub_data(well_list)
+
     def has_incomplete_data(self, col_name: str):
         """
         Checks if a column contains empty cells.

primo/demo/PRIMO - Example_1.ipynb

@@ -1396,7 +1396,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 16,
+   "execution_count": null,
    "id": "27fdb5a7",
    "metadata": {
     "jupyter": {
@@ -1418,11 +1418,9 @@
    "source": [
     "# Use the oil_gdf and gas_gdf pointers defined above\n",
     "cluster_col_name = \"Clusters\"\n",
-    "n_clusters_gas = len(set(gas_wells[cluster_col_name]))  # Number of clusters of gas wells\n",
-    "n_clusters_oil = len(set(oil_wells[cluster_col_name]))  # Number of clusters of oil wells\n",
     "\n",
-    "color_list_oil = get_cluster_colors(n_clusters_oil, pd.unique(oil_wells[cluster_col_name]))\n",
-    "color_list_gas = get_cluster_colors(n_clusters_gas, pd.unique(gas_wells[cluster_col_name]))\n",
+    "color_list_oil = get_cluster_colors(pd.unique(oil_wells[cluster_col_name]))\n",
+    "color_list_gas = get_cluster_colors(pd.unique(gas_wells[cluster_col_name]))\n",
     "\n",
     "# Scatter plot of clusters\n",
     "cluster_figure, (ax1, ax2) = plt.subplots(1, 2, figsize=(20, 8))\n",
@@ -6137,7 +6135,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3",
+   "display_name": "primo-opt",
    "language": "python",
    "name": "python3"
   },
@@ -6151,7 +6149,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.4"
+   "version": "3.12.6"
   }
  },
  "nbformat": 4,

primo/opt_model/result_parser.py

@@ -13,7 +13,7 @@
 
 # Standard libs
 import logging
-from typing import List, Union
+from typing import List, Optional, Union
 
 # Installed libs
 import matplotlib.pyplot as plt
@@ -329,6 +329,25 @@ def __init__(self, wd: WellData, clusters_dict: dict, plugging_cost: dict):
         self.num_projects = len(self.projects)
         self.efficiency_calculator = EfficiencyCalculator(self)
 
+    def get_project_id_by_well_id(self, well_id: str) -> Optional[int]:
+        """
+        Returns the project_id associated with the given well_id.
+
+        Parameters
+        ----------
+        well_id : str
+            The ID of the well.
+
+        Returns
+        -------
+        Optional[int]
+            The project ID if the well exists in any project; otherwise, None.
+        """
+        for project_id, project in self.projects.items():
+            if well_id in project.well_data.data[self.wd.col_names.well_id].values:
+                return project_id
+        return None
+
     def __str__(self) -> str:
         msg = (
             f"The optimal campaign has {self.num_projects} projects.\n"