Files
mcm-mfp/task1/02_neighbor_demand.py

60 lines
1.7 KiB
Python
Raw Normal View History

2026-01-19 01:40:19 +08:00
from __future__ import annotations
import math
import numpy as np
import pandas as pd
INPUT_XLSX = "task1/01_clean.xlsx"
OUTPUT_XLSX = "task1/02_neighbor.xlsx"
RHO_MILES_LIST = [10.0, 20.0, 30.0]
def haversine_miles(lat1: float, lon1: float, lat2: float, lon2: float) -> float:
r_miles = 3958.7613
phi1 = math.radians(lat1)
phi2 = math.radians(lat2)
dphi = math.radians(lat2 - lat1)
dlambda = math.radians(lon2 - lon1)
a = math.sin(dphi / 2.0) ** 2 + math.cos(phi1) * math.cos(phi2) * math.sin(dlambda / 2.0) ** 2
c = 2.0 * math.atan2(math.sqrt(a), math.sqrt(1.0 - a))
return r_miles * c
def main() -> None:
sites = pd.read_excel(INPUT_XLSX, sheet_name="sites")
lat = sites["lat"].to_numpy(float)
lon = sites["lon"].to_numpy(float)
mu = sites["mu_clients_per_visit"].to_numpy(float)
n = len(sites)
dist = np.zeros((n, n), dtype=float)
for i in range(n):
for j in range(i + 1, n):
d = haversine_miles(lat[i], lon[i], lat[j], lon[j])
dist[i, j] = d
dist[j, i] = d
out = sites.copy()
out["neighbor_demand_mu_self"] = mu.copy()
for rho in RHO_MILES_LIST:
w = np.exp(-(dist**2) / (2.0 * rho * rho))
# D_i(rho) = sum_j mu_j * exp(-dist(i,j)^2 / (2 rho^2))
out[f"neighbor_demand_mu_rho_{int(rho)}mi"] = w @ mu
dist_df = pd.DataFrame(dist, columns=sites["site_id"].tolist())
dist_df.insert(0, "site_id", sites["site_id"])
with pd.ExcelWriter(OUTPUT_XLSX, engine="openpyxl") as writer:
out.to_excel(writer, index=False, sheet_name="sites")
dist_df.to_excel(writer, index=False, sheet_name="dist_miles")
if __name__ == "__main__":
main()