Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
A script to calculate demographic gravitational force from a spreadsheet of cities to a given city.
"""Calculates demographic gravitation for a spreadsheet of city data.
This script is specifically designed to work with the basic US cities
spreadsheet at https://simplemaps.com/data/us-cities. However, it could
be made to work with any spreadsheet of cities, latitudes, longitudes,
and populations as long as the sheet name and column names are updated
appropriately.
"""
import pandas as pd
import math
CITY_DATA = "./simplemaps/uscities.xlsx" # City data spreadsheet
HOME_CITY = ("Beavercreek", "OH") # (Home city, home state abbreviation)
DEG_TO_RAD = math.pi / 180
EARTH_MEAN_RADIUS = 3958.7613 # miles
def distance(home_coords, row):
"""Calculates distance between home and a row's city."""
away_coords = (row['lat'], row['lng'])
# Calculate great circle distance using the Haversine formula.
phi_1 = home_coords[0] * DEG_TO_RAD
phi_2 = away_coords[0] * DEG_TO_RAD
delta_phi = ((away_coords[0] - home_coords[0]) * DEG_TO_RAD)
delta_lambda = ((away_coords[1] - home_coords[1]) * DEG_TO_RAD)
a = (math.sin(delta_phi/2)**2
+ math.cos(phi_1)
* math.cos(phi_2)
* math.sin(delta_lambda/2)**2
)
distance = (2
* EARTH_MEAN_RADIUS
* math.atan2(math.sqrt(a), math.sqrt(1-a))
)
return distance
def gravity(home_pop, row):
""" Calculates 'gravitational force' between home and a row's city.
The demographic gravitational force is determined by dividing the
product of the populations of the home city and row city by the
square of their distance.
"""
row_pop = row["population"]
dist = row["distance_mi"]
if dist == 0:
return float("inf")
else:
return (home_pop * row_pop) / (dist ** 2)
# Read city data spreadsheet.
columns = ["city", "state_id", "lat", "lng", "population"]
cities = pd.read_excel(
CITY_DATA, sheet_name="Sheet1", engine="openpyxl"
)[columns]
# Get coordinates of home city.
home = cities.loc[
(cities["city"] == HOME_CITY[0]) & (cities["state_id"] == HOME_CITY[1])
]
home_coords = (home["lat"].values[0], home["lng"].values[0])
home_pop = home["population"].values[0]
# Create distances column.
cities["distance_mi"] = cities.apply(
lambda row: distance(home_coords, row),
axis=1
)
# Remove rows with zero distance:
cities = cities.loc[cities["distance_mi"] != 0]
# Create population gravity column.
cities["pop_grav"] = cities.apply(
lambda row: gravity(home_pop, row),
axis=1
)
# Sort by population gravity.
cities = cities.sort_values(
"pop_grav",
ascending=False,
na_position='first',
ignore_index=True
)
cities.to_csv("./output/pop-grav.csv")
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment