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

import skmob
from skmob.models import gravity
import geopandas as gpd
import matplotlib.pyplot as plt
import pandas as pd
import numpy as np
from sklearn import preprocessing
np.random.seed(0)

station_df = pd.read_csv('./data/station-coor.csv', index_col=0)
station_df.head(2)

data_X_all = gpd.read_file('./data/data_X_all.shp').iloc[:48, :]
tess = data_X_all[['TotPop', 'nearest st', 'geometry']]
print(tess.shape)
tess.head(2)

OD = np.load('./data/3d_daily.npy').sum(axis=2)[:48, :48]
OD_df = pd.DataFrame(OD, index=tess['nearest st'].values, columns=tess['nearest st'].values)

OD_df.shape

ODa = OD_df.stack().reset_index().set_axis(['origin', 'destination', 'flow'], axis=1)
ODa = ODa.sort_values(by=['origin', 'destination']).reset_index(drop=True)
fdf = skmob.FlowDataFrame(ODa, origin='origin', destination='destination', flow='flow', tessellation=tess,
                          tile_id='nearest st')

tot_outflows = fdf[fdf['origin']!=fdf['destination']].groupby(by='origin', axis=0)[['flow']].sum()
tess = tess.merge(tot_outflows, left_on='nearest st', right_on='origin').rename(columns={'flow': 'tot_outflow'})
tess.head()

fdf_train = fdf[:1612]
fdf_train.shape

count = [1 if fdf_train.loc[i, 'origin']==fdf_train.loc[i, 'destination'] else 0 for i in fdf_train.index]
sum(count)

ODb = ODa.pivot(index='origin', columns='destination', values='flow')
ODb.head()

OD_ = np.array(ODb)  # shape: [num_cities, num_cities]
OD_ = OD_.reshape(-1, 1)
scalar_OD = preprocessing.MinMaxScaler().fit(OD_)

plt.imshow(ODb)
plt.xlabel('Origin ID')
plt.ylabel('Destination ID')
plt.title('Real data (2019 sum)', fontsize=11)
plt.colorbar().formatter.set_powerlimits((0, 0))

# # Gravity model

# + [markdown] tags=[]
# ## Singly constrained
# -

# ### Power law deterrence function

singly_fitted1 = gravity.Gravity(gravity_type='singly constrained')
print(singly_fitted1)

singly_fitted1.fit(fdf_train, relevance_column='TotPop')
print(singly_fitted1)

fdf_fit1 = singly_fitted1.generate(tess, tile_id_column='nearest st', tot_outflows_column='tot_outflow',
                                   relevance_column='TotPop', out_format='flows')

fdf_fit1

OD_fit1 = pd.DataFrame(fdf_fit1[['origin', 'destination', 'flow']])
OD_rec1 = OD_fit1.pivot(index='origin', columns='destination', values='flow').fillna(0)
print(OD_rec1.shape)
OD_rec1.head()

plt.imshow(OD_rec1, vmin=0, vmax=9e4)
plt.xlabel('Origin ID')
plt.ylabel('Destination ID')
plt.title('Singly constrained, power law', fontsize=11)
plt.colorbar().formatter.set_powerlimits((0, 0))

# +
# OD_rec1.stack().reset_index().iloc[1612:, :].to_csv('./res/dg and g/g_pred.csv')

StarcoderdataJupyterScriptsDedupFiltered