Changing AREA_DIST_MATRIX to integer keys.

pnxenopoulos · Jul 10, 2023 · 71226ee · 71226ee
1 parent a4f33f2
commit 71226ee
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 33 deletions.
diff --git a/awpy/analytics/nav.py b/awpy/analytics/nav.py
@@ -558,17 +558,17 @@ def generate_area_distance_matrix(map_name: str, *, save: bool = False) -> AreaM
             # Compute center of second area
             area2_x, area2_y, area2_z = _get_area_center(map_name, area2)
             # Calculate basic euclidean distance
-            area_distance_matrix[str(area1)][str(area2)]["euclidean"] = math.sqrt(
+            area_distance_matrix[area1][area2]["euclidean"] = math.sqrt(
                 (area1_x - area2_x) ** 2
                 + (area1_y - area2_y) ** 2
                 + (area1_z - area2_z) ** 2
             )
             # Also get graph distance
-            area_distance_matrix[str(area1)][str(area2)]["graph"] = area_distance(
+            area_distance_matrix[area1][area2]["graph"] = area_distance(
                 map_name, area1, area2, dist_type="graph"
             )["distance"]
             # And geodesic like distance
-            area_distance_matrix[str(area1)][str(area2)]["geodesic"] = area_distance(
+            area_distance_matrix[area1][area2]["geodesic"] = area_distance(
                 map_name, area1, area2, dist_type="geodesic"
             )["distance"]
     if save:
@@ -639,7 +639,7 @@ def _get_median_place_distance(
     connections = []
     for sub_area1 in area_mapping[place1]:
         connections.extend(
-            AREA_DIST_MATRIX[map_name][str(sub_area1)][str(sub_area2)][dist_type]
+            AREA_DIST_MATRIX[map_name][sub_area1][sub_area2][dist_type]
             for sub_area2 in area_mapping[place2]
         )
     return median(connections)
@@ -706,14 +706,10 @@ def generate_place_distance_matrix(map_name: str, *, save: bool = False) -> Plac
                 else:
                     place_distance_matrix[place1][place2][dist_type][
                         "centroid"
-                    ] = AREA_DIST_MATRIX[map_name][str(centroid1)][str(centroid2)][
-                        dist_type
-                    ]
+                    ] = AREA_DIST_MATRIX[map_name][centroid1][centroid2][dist_type]
                     place_distance_matrix[place1][place2][dist_type][
                         "representative_point"
-                    ] = AREA_DIST_MATRIX[map_name][str(reps[place1])][
-                        str(reps[place2])
-                    ][
+                    ] = AREA_DIST_MATRIX[map_name][reps[place1]][reps[place2]][
                         dist_type
                     ]
                     place_distance_matrix[place1][place2][dist_type][
@@ -946,7 +942,7 @@ def _check_arguments_position_distance(
         ValueError: If number of features is not 3 for euclidean distance_type
 
     Returns:
-       tuple[npt.NDArray, npt.NDArray]: Potentially reordered position arrays.
+        tuple[npt.NDArray, npt.NDArray]: Potentially reordered position arrays.
     """
     if map_name not in NAV:
         msg = "Map not found."
@@ -1110,8 +1106,8 @@ def _graph_based_position_distance(
         )
     else:
         this_dist = min(
-            AREA_DIST_MATRIX[map_name][str(area1)][str(area2)][distance_type],
-            AREA_DIST_MATRIX[map_name][str(area2)][str(area1)][distance_type],
+            AREA_DIST_MATRIX[map_name][area1][area2][distance_type],
+            AREA_DIST_MATRIX[map_name][area2][area1][distance_type],
         )
     if this_dist == float("inf"):
         this_dist = sys.maxsize / 6
@@ -1147,10 +1143,10 @@ def position_state_distance(
 
     Raises:
         ValueError: If map_name is not in awpy.data.NAV.
-        ValueError: If distance_type is not one of ["graph", "geodesic", "euclidean"].
-        ValueError: If the 0th (number of teams) and 2nd (number of features) dimensions
+                    If distance_type is not one of ["graph", "geodesic", "euclidean"].
+                    If the 0th (number of teams) and 2nd (number of features) dimensions
                         of the inputs do not have the same size.
-        ValueError: If number of features is not 3 for euclidean distance_type
+                    If number of features is not 3 for euclidean distance_type
     """
     position_array_1, position_array_2 = _check_arguments_position_distance(
         map_name, position_array_1, position_array_2, distance_type

diff --git a/awpy/types.py b/awpy/types.py
@@ -62,7 +62,7 @@ class Area(TypedDict):
 
 
 DistanceType = Literal["graph", "geodesic", "euclidean"]
-AreaMatrix = dict[str, dict[str, dict[DistanceType, float]]]
+AreaMatrix = dict[int, dict[int, dict[DistanceType, float]]]
 PlaceMatrix = dict[
     str,
     dict[

diff --git a/tests/test_nav.py b/tests/test_nav.py
@@ -94,28 +94,28 @@ def setup_class(self):
         )
 
         self.expected_area_matrix = {
-            "1": {
-                "1": {"euclidean": 0, "graph": 0, "geodesic": 0},
-                "2": {"euclidean": 1.0, "graph": 1.0, "geodesic": 1.0},
-                "3": {"euclidean": 2.0, "graph": 1.0, "geodesic": 2.0},
+            1: {
+                1: {"euclidean": 0, "graph": 0, "geodesic": 0},
+                2: {"euclidean": 1.0, "graph": 1.0, "geodesic": 1.0},
+                3: {"euclidean": 2.0, "graph": 1.0, "geodesic": 2.0},
             },
-            "2": {
-                "1": {
+            2: {
+                1: {
                     "euclidean": 1.0,
                     "graph": float("inf"),
                     "geodesic": float("inf"),
                 },
-                "2": {"euclidean": 0, "graph": 0, "geodesic": 0},
-                "3": {
+                2: {"euclidean": 0, "graph": 0, "geodesic": 0},
+                3: {
                     "euclidean": math.sqrt(5),
                     "graph": float("inf"),
                     "geodesic": float("inf"),
                 },
             },
-            "3": {
-                "1": {"euclidean": 2.0, "graph": 1.0, "geodesic": 2.0},
-                "2": {"euclidean": math.sqrt(5), "graph": 2.0, "geodesic": 3.0},
-                "3": {"euclidean": 0, "graph": 0, "geodesic": 0},
+            3: {
+                1: {"euclidean": 2.0, "graph": 1.0, "geodesic": 2.0},
+                2: {"euclidean": math.sqrt(5), "graph": 2.0, "geodesic": 3.0},
+                3: {"euclidean": 0, "graph": 0, "geodesic": 0},
             },
         }
         self.expected_place_matrix_1 = {
@@ -742,11 +742,9 @@ def test_generate_area_distance_matrix(self):
 
         assert isinstance(result_matrix, dict)
         for area1_id in result_matrix:
-            assert isinstance(area1_id, str)
-            assert str(int(area1_id)) == area1_id
+            assert isinstance(area1_id, int)
             for area2_id in result_matrix[area1_id]:
-                assert isinstance(area2_id, str)
-                assert str(int(area2_id)) == area2_id
+                assert isinstance(area2_id, int)
                 for dist_type in result_matrix[area1_id][area2_id]:
                     assert dist_type in {"geodesic", "euclidean", "graph"}
                     assert isinstance(