distance-matrix/distance_matrix.py

#!/usr/bin/python
# -*- coding: utf-8 -*-
# Author: Omar Vega Ramos
# E-mail: ovruni@riseup.net
# License: GNU GPL - GNU General Public License v3.0 or later
# http://www.gnu.org/licenses/gpl.html

from argparse import ArgumentParser
from csv import reader, writer
from math import sin, cos, sqrt, atan2, radians
from os.path import isfile

def calculate_distance(lat1, lon1, lat2, lon2):
    R = 6373.0 # Radius of the earth in km
    lat1 = radians(lat1)
    lon1 = radians(lon1)
    lat2 = radians(lat2)
    lon2 = radians(lon2)

    dlon = lon2 - lon1
    dlat = lat2 - lat1

    a = sin(dlat / 2)**2 + cos(lat1) * cos(lat2) * sin(dlon / 2)**2
    c = 2 * atan2(sqrt(a), sqrt(1 - a))

    distance = R * c # Distance in km

    return distance

parser = ArgumentParser(description='Calculate the distances between the coordinates of two CSV files.')
parser.add_argument('input1', help='name of the first input file')
parser.add_argument('input2', help='name of the second input file')
parser.add_argument('--lat', default='lat', help='column name for latitude')
parser.add_argument('--lon', default='lon', help='column name for longitude')
parser.add_argument('-o', '--output', default='output.csv', help='name of the output file')
args = parser.parse_args()

if isfile(args.input1) == False:
    exit("{}: error: the file '{}' does not exist".format(parser.prog, args.input1))

if isfile(args.input2) == False:
    exit("{}: error: the file '{}' does not exist".format(parser.prog, args.input2))

table1 = []
table2 = []
header1 = None
header2 = None
lat_pos1 = None
lon_pos1 = None
lat_pos2 = None
lon_pos2 = None

with open(args.input1) as csv_file:
    csv_reader = reader(csv_file, delimiter=',')
    is_header = True

    for row in csv_reader:
        if is_header == True:
            header1 = row

            for i, col in enumerate(header1):
                if col == args.lat:
                    lat_pos1 = i
                elif col == args.lon:
                    lon_pos1 = i

            if lat_pos1 == None:
                print("ERROR: Column ´{}´ does not exist in file ´{}´".format(args.lat, args.input1))
                quit()
            if lon_pos1 == None:
                print("ERROR: Column ´{}´ does not exist in file ´{}´".format(args.lon, args.input1))
                quit()

            is_header = False
            continue

        table1.append(row)

with open(args.input2) as csv_file:
    csv_reader = reader(csv_file, delimiter=',')
    is_header = True

    for row in csv_reader:
        if is_header == True:
            header2 = row
            for i, col in enumerate(header2):
                if col == args.lat:
                    lat_pos2 = i
                elif col == args.lon:
                    lon_pos2 = i

            if lat_pos2 == None:
                print("ERROR: Column ´{}´ does not exist in file ´{}´".format(args.lat, args.input2))
                quit()
            if lon_pos2 == None:
                print("ERROR: Column ´{}´ does not exist in file ´{}´".format(args.lon, args.input2))
                quit()
            is_header = False
            continue

        table2.append(row)

rows = []
row = header1 + header2 + ['distance (km)']
rows.append(row)

for row1 in table1:
    lat1 = row1[lat_pos1]
    lon1 = row1[lon_pos1]

    if lat1 == '' or lon1 == '':
        print("ERROR: Wrong coordinates in {} file.".format(args.input1))
        quit()

    lat1 = float(lat1)
    lon1 = float(lon1)

    for row2 in table2:
        lat2 = row2[lat_pos2]
        lon2 = row2[lon_pos2]

        if lat2 == '' or lon2 == '':
            print("ERROR: Wrong coordinates in {} file.".format(args.input2))
            quit()

        lat2 = float(lat2)
        lon2 = float(lon2)
        distance = calculate_distance(lat1, lon1, lat2, lon2)
        row = row1 + row2 + [distance]
        rows.append(row)

with open(args.output, 'w', encoding='UTF8', newline='') as f:
    writer = writer(f)
    writer.writerows(rows)

print("The ´{}´ file was generated.".format(args.output))