# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # Earthquake localization using progressive search grid
#
# Using pytorch GPU-accelerated engine

# !pip install python-dateutil geopy

# +
from datetime import datetime
import dateutil.parser

# panyaungan2018
eq_lat = -7.092
eq_lon = 105.963
eq_epicenter = (eq_lat, eq_lon)
eq_depth = 48.2
eq_mw = 6.0
eq_origin_time = dateutil.parser.parse('2018-01-23T06:34:54.000+00:00')

print('EQ origin_time=%s' % (eq_origin_time))


# +
import pandas as pd

# arrivals = pd.read_csv('panyaungan2018_real.csv')
arrivals = pd.read_csv('panyaungan2018_augmented.csv')
arrivals
# -

# Parse arrival time
# arrivals['arrivalTime_dt'] = arrivals.apply(lambda row: dateutil.parser.parse(row['arrivalTime']), axis=1)
arrivals['arrivalTime_ts'] = arrivals.apply(lambda row: dateutil.parser.parse(row['arrivalTime']).timestamp(), axis=1)
arrivals['latency'] = arrivals['arrivalTime_ts'] - eq_origin_time.timestamp()
arrivals

# Sort based on earliest arrival time
arrivals_sorted = arrivals.sort_values('arrivalTime_ts')
arrivals_sorted

# Cut only 3 nearest stations
arrivals_sorted = arrivals_sorted[:3]
arrivals_sorted

# Earliest station
p1 = arrivals_sorted.iloc[1]
p1

# Arrival time of earliest station
t1 = p1['arrivalTime_ts']
t1

p1_loc = (p1['lat'], p1['lon'])
p1_loc


# +
# Make a grid using 3-deg radius
def round_res(num: float, resolution: float): 
    return round(num / resolution) * resolution

p1_locr = (round_res(p1_loc[0], 0.3), round_res(p1_loc[1], 0.3))
p1_locr

# +
import numpy as np
import torch
import torch.cuda

point_interval = 0.1
lats = torch.arange(p1_locr[0] - 1.2, p1_locr[0] + 1.2 + 0.001, point_interval, dtype=torch.float, device='cuda')
lons = torch.arange(p1_locr[1] - 1.2, p1_locr[1] + 1.2 + 0.001, point_interval, dtype=torch.float, device='cuda')
n_points = len(lats) * len(lons)
print('lats=%s lons=%s len=%dx%d n_points=%d' % (lats, lons, len(lats), len(lons), n_points))

# -

grid = torch.empty([n_points, 2], dtype=torch.float, device='cuda') # our lat lon is just approximate
grid[0]

for lat_idx, lat in enumerate(lats):
    for lon_idx, lon in enumerate(lons):
        idx = lat_idx * len(lons) + lon_idx
        grid[idx, 0] = lat
        grid[idx, 1] = lon
grid[:5]

# +
from datetime import timedelta

times = torch.empty(31, dtype=torch.float64, device='cuda') # MUST be 64-bit!
for idx, sec in enumerate(range(31)):
    print('%s - %s = %s' % (t1, sec, t1 - sec))
    times[idx] = t1 - sec
print(len(times))
times

# +
eq_vel_deg = 0.067 # Primary wave velocity = 0.067 degree / second
eq_vel_m = 7437 # Primary wave velocity = 0.067 deg/s * 111000 = 7437 m/s

velocities = torch.empty(1, dtype=torch.float16, device='cuda')
velocities[0] = eq_vel_m
velocities
# -

# ## Weighted RMSE
#
# $$\text{wRMSE} = \sqrt{\sum_{i=1}^n w_i (\hat T_i - T_i)^2}$$
#

from geopy import distance

# +
# https://gist.github.com/nickjevershed/6480846
import math

def distance_on_sphere(a: torch.Tensor, b: torch.Tensor, radius: float):

    # Convert latitude and longitude to 
    # spherical coordinates in radians.
    degrees_to_radians = math.pi/180.0
        
    # phi = 90 - latitude
    phi1 = (90.0 - a.select(4, 0)) * degrees_to_radians
    phi2 = (90.0 - b.select(4, 0)) * degrees_to_radians
        
    # theta = longitude
    theta1 = a.select(4, 1) * degrees_to_radians
    theta2 = b.select(4, 1) * degrees_to_radians
        
    # Compute spherical distance from spherical coordinates.
        
    # For two locations in spherical coordinates 
    # (1, theta, phi) and (1, theta, phi)
    # cosine( arc length ) = 
    #    sin phi sin phi' cos(theta-theta') + cos phi cos phi'
    # distance = rho * arc length
    
    cos = (torch.sin(phi1) * torch.sin(phi2) * torch.cos(theta1 - theta2) + 
           torch.cos(phi1) * torch.cos(phi2))
    arc = torch.acos( cos )

    # Remember to multiply arc by the radius of the earth 
    # in your favorite set of units to get length.
    return arc * radius



# +
def score_stations(station_locs: torch.Tensor, actual_arrivals: torch.Tensor,
                  epicenter_es: torch.Tensor, origin_time_es: torch.Tensor, velocity_es: torch.Tensor):
    '''Calculate distance (m) and the error of arrival time (seconds) for given stations and given solution parameters (epicenter, origin time, velocity).'''
    # Returned shape is E x O x V x S x 2
    # where E=length of epicenters, O=length of origin times, V=length of velocities, S=length of stations

#     print('Dists traditional...')
#     dists = torch.empty(len(epicenter_es), len(origin_time_es), len(velocity_es), len(station_locs), 
#                         dtype=torch.float64, device='cuda')
#     for epicenter_idx, epicenter_e in enumerate(epicenter_es):
#         for station_idx, station_loc in enumerate(station_locs):
#             dist = distance.distance(epicenter_e, station_loc).m
#             for origin_time_idx, origin_time_e in enumerate(origin_time_es):
#                 for velocity_idx, velocity_e in enumerate(velocity_es):
#                     dists[epicenter_idx][origin_time_idx][velocity_idx][station_idx] = dist

#     print('Dists geodesic + expand...')
#     dists = torch.empty(len(epicenter_es), 1, 1, len(station_locs), 
#                         dtype=torch.float64, device='cuda')
# #     print('dists: %s' % (dists.size(),))
#     for epicenter_idx, epicenter_e in enumerate(epicenter_es):
#         for station_idx, station_loc in enumerate(station_locs):
#             dist = distance.distance(epicenter_e, station_loc).m
# #             dist = distance.great_circle(epicenter_e, station_loc).m
#             dists[epicenter_idx][0][0][station_idx] = dist
#     print('Dists expand...')
#     dists.expand(-1, len(origin_time_es), len(velocity_es), -1)
    
    print('Dists great circle (GPU) + expand...')
    dists1 = torch.empty(len(epicenter_es), 1, 1, len(station_locs), 2,
                        dtype=torch.float64, device='cuda')
    dists2 = torch.empty(len(epicenter_es), 1, 1, len(station_locs), 2,
                        dtype=torch.float64, device='cuda')
    for epicenter_idx, epicenter_e in enumerate(epicenter_es):
        for station_idx, station_loc in enumerate(station_locs):
            dists1[epicenter_idx][0][0][station_idx] = epicenter_e
            dists2[epicenter_idx][0][0][station_idx] = station_loc
    dists = distance_on_sphere(dists1, dists2, distance.EARTH_RADIUS * 1000)
    print('Dists expand...')
    dists.expand(-1, len(origin_time_es), len(velocity_es), -1)
        
#     print('dists = %s' % dists)
    print('Travel durations...')
    travel_durs = dists / eq_vel_m
#     print('travel_durs %s = %s' % (travel_durs.size(), travel_durs))
    print('Arrival times...')
    origin_time_x = origin_time_es.view(1, len(origin_time_es), 1, 1).expand(len(epicenter_es), -1, len(velocity_es), len(station_locs))
    arrival_times_es = origin_time_x + travel_durs
#     print('arrival_times_es = %s' % arrival_times_es)
#     print('actual_arrivals = %s' % actual_arrivals)
    print('Arrival time errors...')
    arrival_errors = arrival_times_es - actual_arrivals
#     print('arrival_errors = %s' % arrival_errors)
#     print('From epicenter %s to station %s %s (%d km) in %s s' % 
#           (epicenter_e, station_code, station_loc, dist_m/1000, travel_dur))
#     print('Est arrival %s actual %s error=%s' % 
#           (arrival_time_e, actual_arrival, arrival_error))
    return (dists, arrival_errors)

station_codes = list(arrivals_sorted['station'])
print(station_codes)
station_locs = torch.Tensor([arrivals_sorted['lat'], arrivals_sorted['lon']]).transpose(0, 1)
print(station_locs)
actual_arrivals = torch.cuda.DoubleTensor(arrivals_sorted['arrivalTime_ts'])
print(actual_arrivals)

epicenter_e = torch.tensor([-9.899999999999999, 103.5], dtype=torch.float, device='cuda')
origin_time_e = dateutil.parser.parse('2018-01-23T06:35:03.851000+00:00').timestamp()
print('epicenter_e = %s' % epicenter_e)
epicenter_es = torch.unsqueeze(epicenter_e, 0)
print('epicenter_es = %s' % epicenter_es)
origin_time_es = torch.tensor([origin_time_e], dtype=torch.double, device='cuda')
print(origin_time_es)
score_stations(station_locs, actual_arrivals, epicenter_es, origin_time_es, velocities)
# -

# Estimated wrong one with 0.0020 RWSSE
epicenter_es = torch.unsqueeze(grid[111], 0)
origin_time_es = torch.unsqueeze(times[28], 0)
score_stations(station_locs, actual_arrivals, epicenter_es, origin_time_es, velocities)

# Should be correct one? with 0.0020 RWSSE
grid
# epicenter_es = torch.unsqueeze(grid[111], 0)
# origin_time_es = torch.unsqueeze(times[28], 0)
#score_stations(station_locs, actual_arrivals, epicenter_es, origin_time_es, velocities)

# +
# MULTI-SOLUTION ROWS VERSION
def score_solutions(station_codes: list, station_locs: torch.Tensor, actual_arrivals: torch.Tensor,
          epicenter_es: torch.Tensor, origin_time_es: torch.Tensor, velocity_es: torch.Tensor):
    '''Returns the RWSSE for all stations given parameters containing estimates.'''
    
#     score_station(stations.iloc[0]['station'], (stations.iloc[0]['lat'], stations.iloc[0]['lon']),
#                   stations.iloc[0]['arrivalTime_dt'],
#                   epicenter_e, origin_time_e)
    (dists, arrival_errors) = score_stations(station_locs, actual_arrivals,
                             epicenter_es, origin_time_es, velocity_es)
#     print(result_arr)
    print('Square error...')
    ses = torch.pow(arrival_errors, 2)
    print('Weighted SE...')
    weighted_ses = torch.mul(1 / dists, ses)
#     weighted_ses = ses
#     print(weighted_ses)
    print('Sum...')
    wsse = weighted_ses.sum(3, dtype=torch.float16) # only need for reasonable precision here
#     print(wsse)
    print('Root...')
    rwsse = torch.sqrt(wsse)
#     print('RWSSEs: %s' % rwsse)
    return rwsse

station_codes = list(arrivals_sorted['station'])
# print(station_codes)
station_locs = torch.cuda.HalfTensor([arrivals_sorted['lat'], arrivals_sorted['lon']]).transpose(0, 1)
#print(station_locs)
actual_arrivals = torch.cuda.DoubleTensor(arrivals_sorted['arrivalTime_ts'])
# print(actual_arrivals)
# print(grid[:2])
# print(velocities)
# print(times[:15])
score_solutions(station_codes, station_locs, actual_arrivals, grid[:2], times[:2], velocities)

# +
# Let's score everything!
import time

print('Scoring %d x %d x %d = %d combinations...' % 
      (len(grid), len(times), len(velocities), len(grid) * len(times) * len(velocities)))
start = time.process_time()
scores = score_solutions(station_codes, station_locs, actual_arrivals, grid, times, velocities)
elapsed = time.process_time() - start
print('Scored in %s seconds' % elapsed)
print(scores.size())
# -

scores

best_score = torch.min(scores)
best_score

best_idxs = (scores == best_score).nonzero()
best_idx = best_idxs[0]
best_idx

times[28].item()

(epicenter_e_idx, origin_time_e_idx, velocity_e_idx) = best_idx
epicenter_e_idx

# +
from datetime import timezone

epicenter_e = grid[epicenter_e_idx]
origin_time_e = datetime(1970, 1, 1, tzinfo=timezone.utc) + timedelta(seconds=times[origin_time_e_idx].item())
velocity_e = velocities[velocity_e_idx]
# -

t_last_offset = np.round((arrivals_sorted.tail(1)['arrivalTime_ts'].item() - eq_origin_time.timestamp()) / 0.001) * 0.001
epicenter_err = distance.distance(epicenter_e, eq_epicenter).km
origin_time_err = abs((origin_time_e - eq_origin_time).total_seconds())
print('Estimated epicenter: %s. Actual epicenter: %s.' % (np.array(epicenter_e), eq_epicenter))
print('Estimated origin time: %s. Actual origin time: %s.' % (origin_time_e, eq_origin_time))
print('Estimated P-wave velocity: %s m/s' % velocity_e.item())
print('Decision latency: T + %s seconds. Epicenter error: ±%d km. Origin time error: ±%s seconds.' %
      (t_last_offset, epicenter_err, origin_time_err))



StarcoderdataJupyterScriptsDedupFiltered