A script to calculate demographic gravitational force from a spreadsheet of cities to a given city.

population-gravity.py

"""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")