citybert/generate_data.py

import concurrent.futures
import csv
import itertools
from concurrent.futures import as_completed
from functools import lru_cache

import geonamescache
import numpy as np
from geopy.distance import geodesic

MAX_DISTANCE = 20_037.5

gc = geonamescache.GeonamesCache()
cities = gc.get_cities()
us_cities = {
    k: c
    for k, c in cities.items()
    if (c.get("countrycode") == "US")  # & (c.get("population", 0) > 5e4)
}


@lru_cache(maxsize=None)
def get_coordinates(city_name, country_code="US"):
    """
    Get the coordinates of a city.

    Parameters
    ----------
    city_name : str
        The name of the city.
    country_code : str, optional
        The country code of the city, by default 'US'.

    Returns
    -------
    tuple
        A tuple containing the latitude and longitude of the city,
        or None if the city is not found.
    """
    search_results = gc.search_cities(city_name, case_sensitive=True)
    search_results = {
        k: c
        for k, c in search_results.items()
        if (c.get("countrycode") == country_code)
    }

    if not search_results:
        return None
    populations = [city.get("population") for city in search_results]
    city = search_results[np.argmax(populations)]
    return city.get("latitude"), city.get("longitude")


def get_distance(city1, city2, country1="US", country2="US"):
    """
    Get the distance between two cities in kilometers.

    Parameters
    ----------
    city1 : str
        The name of the first city.
    city2 : str
        The name of the second city.
    country1 : str, optional
        The country code of the first city, by default 'US'.
    country2 : str, optional
        The country code of the second city, by default 'US'.

    Returns
    -------
    float
        The distance between the two cities in kilometers,
        or None if one or both city names were not found.
    """
    city1_coords = get_coordinates(city1, country1)
    city2_coords = get_coordinates(city2, country2)

    if (city1_coords is None) or (city2_coords is None):
        return None

    return geodesic(city1_coords, city2_coords).km


def calculate_distance(pair):
    city1, city2 = pair
    distance = get_distance(city1["name"], city2["name"])
    return city1["name"], city2["name"], distance


def main():
    cities = list(us_cities.values())
    print(f"Num cities: {len(cities)}")
    city_combinations = list(itertools.combinations(cities, 2))

    with open("city_distances_full.csv", "w", newline="") as csvfile:
        fieldnames = ["city_from", "city_to", "distance"]
        writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
        writer.writeheader()

        try:
            executor = concurrent.futures.ProcessPoolExecutor(max_workers=8)
            # results = executor.map(calculate_distance, city_combinations)
            futures = {
                executor.submit(calculate_distance, pair): pair
                for pair in city_combinations
            }
            for future in as_completed(futures):
                city_from, city_to, distance = future.result()
                if distance is not None:
                    writer.writerow(
                        {
                            "city_from": city_from,
                            "city_to": city_to,
                            "distance": distance,
                        }
                    )
        except KeyboardInterrupt:
            print("Interrupted. Terminating processes...")
            executor.shutdown(wait=False)
            raise SystemExit("Execution terminated by user.")


if __name__ == "__main__":
    main()