Update utils.py

utils.py

@@ -1,13 +1,31 @@
 from collections import Counter
 import pandas as pd
 import numpy as np
+from scipy.spatial import cKDTree
 
+df_banks = pd.read_csv("df_amenities.csv")
+df_banks = pd.read_csv("df_banks.csv")
+
+bank_coords = df_banks[['lat','lon']].values
+tree_banks = cKDTree(bank_coords)
+
+amenity_coords = df_amenities[['lat','lon']].values
+tree_amenities = cKDTree(amenity_coords)
+
+DATASET_COLUMNS = [
+    'num_Dining and Drinking', 'num_Community and Government', 'num_Retail',
+       'num_Business and Professional Services', 'num_Landmarks and Outdoors',
+       'num_Arts and Entertainment', 'num_Health and Medicine',
+       'num_Travel and Transportation', 'num_Sports and Recreation',
+       'num_Event'
+]
 
 def compute_features(candidate_point, radius=0.005):
     lat, lon = candidate_point
 
     # Banks
     bank_idxs = tree_banks.query_ball_point([lat, lon], r=radius)
+    
     n_banks = len(bank_idxs)
     if n_banks > 0:
         neighbors = df_banks.iloc[bank_idxs]
@@ -37,9 +55,11 @@ def compute_features(candidate_point, radius=0.005):
 
     # Count occurrences per category
     count_per_category = Counter(all_category_ids)
-    for cat, cnt in count_per_category.items():
-        features[f'num_{cat}'] = cnt
+    for feat in DATASET_COLUMNS:
+    # for cat, cnt in count_per_category.items():
+        features[f'num_{feat}'] = count_per_category.get(feat, 0)
 
+    
     # # Count occurrences of first category
     # first_categories = [cats[0] for cats in amenities['fsq_category_ids'] if len(cats)>0]
     # count_first_category = Counter(first_categories)