london_dataset.py

'''London bike sharing dataset class for Pytorch geometric. The class builds on the Dataset parent from
Pytorch Geometric.

Written by: Anders Ohrn, May 2020

'''
import torch
from torch_geometric.data import Dataset, DataLoader
from torch_geometric.data import Data

from os import listdir, makedirs, remove
import os.path as osp
import pandas as pd
import numpy as np
from numpy.random import Generator, PCG64

from consts import EXCLUDE_STATIONS

'''File prefix for the processed files'''
TIME_SLICE_NAME = 'time_slice'

class LondonBikeDataset(Dataset):
    '''The Geometric PyTorch dataset class to make raw data available for PyTorch processing, including DataLoaders

    The workflow is as this:

    1. Data as generated by the `rawdata_reformat` is taken as input.
    2. Contiguous slices of set size of the two data channels are created, defining the input to the algorithm,
       and one point a set number of steps into the future is taken, defining the dependent variable for the algorithm.
    3. Each slice is situated in the graph and pickled, such that a DataLoader rapidly can create batches from it.

    The slicing of the data into input-output pairs for supervised learning can take considerable time (minutes to hours)
    if the range of time indices is wide (1000+). Since the time slices are stored, additional
    training can be done on the same data without rerunning the workflow above. If `create_from_source` is set to `False` the
    prior data is provided to the DataLoader (or other equivalent method). Note that no validation is done to ensure
    the data slices are consistent with the initialization parameters. Set `create_from_source` to `False` with caution.

    Args:
        root_dir (str): Path to directory where to the pickled pre-processed time slices are stored. If `process`
            is `True` the content in this directory is created anew. If `process` is `False` the content in this
            directory is only referenced. The files in this directory are not meant to be moved, read or adjusted.
        create_from_source (bool, optional): Flag to determine if the class initialization should create all the time
            slices from the source or only link to already processed data in the `root_dir`. Default is `True`.
        source_dir (str, optional): Path to source directory for input data to process. This folder has to conform to the
            folder hierarchy defined in class documentation. Typically the data content is created by a prior
            execution of `rawdata_reformat`.
        source_data_files (str, optional): Either the name of the file with the raw data located in `source_dir`, or
            a list of names of the files with the raw data located in `source_dir`.
        source_graph_file (str, optional): File name to the weighted graph that is stored in the source directory and
            which applies to all raw data variants. Default is `graph_weight.csv`
        station_exclusion (list, optional): If certain stations should be entirely excluded from consideration,
            provide their station IDs in a list. Default is no exclusion.
        weight_type (str, optional): The type of graph weight on which to base filtering. Options governed by the
            `rawdata_reformat` and includes `percent_flow`, `median_norm` and `total`.
        lower_weight (float, optional): If the edges in the station-to-station graph should be pruned on basis
            of weight, provide the lower bound of weight for edges to keep. Default is no pruning.
        symmetric (bool, optional): TBD
        time_id_bounds (tuple, optional): If only data for subset of times are to be considered, provide the lower
            and upper bound time indeces as a tuple of two integers. Default is no time slicing.
        time_shuffle (bool, optional): If only a random subset of time slices are to be created, set this to True.
            Default is False.
        sample_size (int, optional): If only a random subset of time slices are to be created, set this to the
            size of the random sample. Default is None.
        stride (int, optional): If only an evenly spaced subetset of time slices are to be created, set this to the
            stride to iterate through the data. Default is None.
        ntimes_leading(int, optional): How many contiguous time intervals of data to provide as input, the "x", to
            the model. Value must be compatible with the model architecture. Default 9.
        ntimes_forward(int, optional): How many time intervals into the future from the most recent given
            time interval given to the prediction that should be predicted by the model, the "y". Default 1.
        common_weight (float, optional): If the weight of the edges, after optional filtering by `lower_weight, should be
            substituted for a common value, provide common value here. Note that if the common weight is set to 1.0,
            the weight matrix will be functionality identical to providing no weight matrix, only adjacency matrix, to
            the Pytorch geometric graph convolution methods. Default is no substitution of weight values.
        seed (int, optional): Random seed for random number generator if time shuffle applies. Default no hard-coded seed.

    '''
    def __init__(self, root_dir,
                 create_from_source=True,
                 source_dir=None, source_data_files=None, source_graph_file='graph_weight.csv',
                 station_exclusion=None, weight_type=None, lower_weight=None, symmetric=True,
                 time_id_bounds=None, time_shuffle=False, sample_size=None, stride=None,
                 ntimes_leading=9, ntimes_forward=1, common_weight=None, seed=None):

        self.root = root_dir

        if create_from_source:
            if any([source_dir, source_data_files, source_graph_file]) is None:
                raise ValueError('To process the dataset, source directory, source files and source graph file all required')
        self.source_dir = source_dir
        self.source_graph_file = source_graph_file
        if isinstance(source_data_files, str):
            self.source_data_files = [source_data_files]

        self.time_input_number = ntimes_leading
        self.time_forward_pred = ntimes_forward
        self.common_weight = common_weight

        self.station_exclusion = station_exclusion
        self.weight_type = weight_type
        self.lower_weight = lower_weight
        self.symmetric = symmetric
        self.time_id_bounds = time_id_bounds
        self.t_shuffle = time_shuffle
        self.sample_size = sample_size
        self.stride = stride

        if seed is None:
            self.rg = Generator(PCG64())
        else:
            self.rg = Generator(PCG64(seed))

        # Either create files or simply link to existing ones. This slightly convoluted way is required because
        # the total number of processed files is only known after the processing is done. Therefore the
        # process command must be made explicit unlike the template Dataset
        if create_from_source:
            if not osp.exists(root_dir):
                makedirs(root_dir)
            else:
                for ff in listdir(root_dir):
                    remove(root_dir + '/' + ff)
            self.create_torch_data()
        else:
            self._processed_file_names = [fname for fname in listdir(root_dir) if TIME_SLICE_NAME in fname]

        super(LondonBikeDataset, self).__init__(root_dir)

    @property
    def raw_file_names(self):
        return self.source_data_files

    @property
    def processed_file_names(self):
        return self._processed_file_names

    def len(self):
        return len(self.processed_file_names)

    def get(self, idx):
        data = torch.load(self.root +'/' + TIME_SLICE_NAME + '_{}.pt'.format(idx))
        return data

    def write_creation_params(self, file_name='params.csv', directory=None):
        '''Write the value of the parameters that can affect the time slices in the root directory to a file.

        '''
        if directory is None:
            dir = self.root
        else:
            dir = directory

        with open(dir + '/' + file_name, 'w') as fin:
            for prop_name in ['source_dir', 'source_data_files', 'source_graph_file', 'time_input_number',
                              'time_forward_pred', 'common_weight', 'station_exclusion', 'lower_weight',
                              'time_id_bounds', 't_shuffle', 'sample_size', 'stride', 'weight_type', 'rg']:
                print('{}, "{}"'.format(prop_name, getattr(self, prop_name)), file=fin)

    def create_torch_data(self):
        '''Method to construct the time slices given the raw data, and populate the weighted graph,
        and then to write the graph to the root directory from which the DataLoader later takes data

        '''
        edge_index, edge_attr = self._make_edges()

        self._processed_file_names = []
        for raw_file_name in self.raw_file_names:
            df = pd.read_csv(self.source_dir + '/' + raw_file_name)

            # Apply filters on data to situate in the graph
            if not self.station_exclusion is None:
                df = df.loc[~df['station_id'].isin(self.station_exclusion)]

            if not self.time_id_bounds is None:
                df = df.loc[(df['time_id'] >= self.time_id_bounds[0]) & \
                            (df['time_id'] <= self.time_id_bounds[1])]

            # In the event that filters removed all data, jump to next iteration
            if len(df) == 0:
                continue

            slice_generator = self._time_windows(df,
                                                 t_start=df['time_id'].min(),
                                                 t_end=df['time_id'].max() - self.time_input_number - self.time_forward_pred)

            for count, (data_g_xt, data_g_yt) in enumerate(slice_generator):
                data_graph = Data(x=torch.tensor(data_g_xt, dtype=torch.float),
                                  y=torch.tensor(data_g_yt, dtype=torch.float),
                                  edge_index=edge_index, edge_attr=edge_attr)

                torch.save(data_graph, self.root + '/' + TIME_SLICE_NAME + '_{}.pt'.format(count))
                self._processed_file_names.append('{}_{}.pt'.format(TIME_SLICE_NAME, count))

    def _time_windows(self, df, t_start, t_end):
        '''Generator to slice data into independent and dependent data

        '''
        if self.sample_size is None and self.t_shuffle is False:
            if self.stride is None:
                t_val_iter = range(t_start, t_end)
            else:
                t_val_iter = range(t_start, t_end, self.stride)

        elif (not self.sample_size is None) and self.t_shuffle:
            if self.stride is None:
                t_val_iter = self.rg.choice(range(t_start, t_end), size=self.sample_size, replace=False)
            else:
                raise ValueError('Time windows cannot be generated with both stride and shuffle')

        elif (not self.sample_size is None) and self.t_shuffle is False:
            raise ValueError('To control sample size without shuffle, set t_start and t_end to appropriate values')

        else:
            raise RuntimeError('Generator parameters completely off!')

        # The generator loop with loop variable the lowest time index for the dependent variables, or the least
        # current data point
        null_list = [0] * self.time_input_number
        for t_val in t_val_iter:
            df_x = df.loc[(df['time_id'] < t_val + self.time_input_number) & (df['time_id'] >= t_val)]
            df_y = df.loc[df['time_id'] == t_val + self.time_input_number + self.time_forward_pred - 1]

            if len(df_x) == 0 or len(df_y) == 0:
                raise RuntimeError('Encountered missing time for time_X: {}, and time_Y: {}'.format(t_val, t_val + self.time_input_number + self.time_forward_pred - 1))

            # Move from pandas to torch tensor compatible data. Because a given time slice may have been created
            # from raw data that does not include events at all stations in the graph, all stations are accounted
            # for through null initialization, which are replaced if actual station data is available.
            data_xt = [[null_list, null_list]] * len(self.station_id_2_node_id_map)
            for station_id, chunk in df_x.groupby('station_id'):
                ind = self.station_id_2_node_id_map[station_id]
                data_xt[ind] = [chunk['arrivals'].tolist(),
                                chunk['departures'].tolist()]

            data_yt = [[0, 0]] * len(self.station_id_2_node_id_map)
            for station_id, chunk in df_y.groupby('station_id'):
                ind = self.station_id_2_node_id_map[station_id]
                data_yt[ind] = [chunk['arrivals'].item(), chunk['departures'].item()]

            yield data_xt, data_yt

    def _make_edges(self):
        '''Convert the graph weight file into edge data compatible with Pytorch geometric conventions

        '''
        df_gw = pd.read_csv('{}/{}'.format(self.source_dir, self.source_graph_file))
        df_gw['StartStation Id'] = df_gw['StartStation Id'].astype(int)
        df_gw['EndStation Id'] = df_gw['EndStation Id'].astype(int)
        df_gw['weight'] = df_gw[self.weight_type].astype(float)
        df_gw = df_gw[['StartStation Id', 'EndStation Id', 'weight']]

        if not self.station_exclusion is None:
            df_gw = df_gw.loc[~df_gw['StartStation Id'].isin(self.station_exclusion)]
            df_gw = df_gw.loc[~df_gw['EndStation Id'].isin(self.station_exclusion)]

        # Create map between station ids and node id, needed because Pytorch geometric reindex nodes
        all_ids = set(df_gw['StartStation Id'].unique()).union(set(df_gw['EndStation Id'].unique()))
        self.node_id_2_station_id_map = dict((n, sid) for n, sid in enumerate(all_ids))
        self.station_id_2_node_id_map = {v: k for k, v in self.node_id_2_station_id_map.items()}

        # Make graph non-directional with averaged weights and without self-loops
        if self.symmetric:
            df1 = df_gw.loc[df_gw['StartStation Id'] != df_gw['EndStation Id']].set_index(['StartStation Id', 'EndStation Id'])
            df2 = df1.reorder_levels(['EndStation Id', 'StartStation Id'])
            df2 = df2.reindex(df1.index, fill_value=0.0)
            df3 = pd.concat([df1, df2], axis=1).mean(axis=1)
            df3 = df3.reset_index()
            df3 = df3.rename(columns={0 : 'weight'})
        else:
            df3 = df_gw

        if not self.lower_weight is None:
            df3 = df3.loc[df3['weight'] >= self.lower_weight]

        if not self.common_weight is None:
            df3['weight'] = self.common_weight

        # Put data in format expected by Pytorch
        index = torch.tensor([[self.station_id_2_node_id_map[k] for k in df3['StartStation Id']],
                              [self.station_id_2_node_id_map[k] for k in df3['EndStation Id']]], dtype=torch.long)
        attr = torch.tensor(df3['weight'].tolist(), dtype=torch.float)

        return index, attr

def test():

    bike_dataset = LondonBikeDataset(root_dir='/Users/andersohrn/PycharmProjects/torch/data_15m_4forward_all_asymw',
                                     source_dir='/Users/andersohrn/Development/london_bike_forecast/data_reformat_May21/1701_2004_15m',
                                     source_data_files='dataraw_15m.csv',
                                     source_graph_file='graph_weight.csv',
                                     weight_type='percent_flow',
                                     common_weight=None,
                                     lower_weight=None,
                                     symmetric=False,
                                     time_shuffle=True,
                                     sample_size=5888,
                                     create_from_source=True,
                                     ntimes_leading=9,
                                     ntimes_forward=4,
                                     station_exclusion=EXCLUDE_STATIONS,
                                     time_id_bounds=(2976, 38015),
                                     #time_id_bounds=(38016, 73055),
                                     seed=810828)
    bike_dataset.write_creation_params()

#    bike_dataset_2 = LondonBikeDataset(root_dir='/Users/andersohrn/PycharmProjects/torch/data_15m_1percent', create_from_source=False)
#
#    for k in DataLoader(bike_dataset_2):
#        print (k)

if __name__ == '__main__':
    test()