station_explorer.py

import pickle
import numpy as np
import pandas as pd
data = pd.read_pickle('dataframe.pkl')

class StationExplorer :

    def __init__(self):
        self.data = data

    def get_km(self,station,line):
        """
        get the kilometer number of a station on the line.

        Parameters:
        station (str): The name of the station.
        line (double): The line number.

        Returns:
        float: The kilometer number of the station on the line.
        """
        return self.data.loc[(self.data['Stop name'] == station) & (self.data['Line'] == line), 'KM'].iloc[0]

    def served_stations(self,line):
        """
        get the stations of a line for a given line number

        Parameters:
        line (double): The line number.

        Returns:
        served_stations (array of strings): The stations of the line that are served by the trainline.
        """
        served_stations = self.data[self.data['Line'] == line]['Stop name'].unique()
        return served_stations
    
    def ordered_served_stations(self,line):
        """
        order the served stations of a line

        Parameters:
        line (double): The line number.

        Returns:
        ordered_served_stations (array of strings): The stations of the line in order of being served by the trainline.
        """
        served_stations = self.served_stations(line)
        ordered_served_stations = sorted(served_stations, key = lambda station: self.get_km(station, line))
        return ordered_served_stations
    
explorer = StationExplorer()

'''
############### DESCRIPTIVE INFORMATION ############################
print(pd.Categorical(data['Station abbreviation']).categories) #811 Stations
print(pd.Categorical(data['Stop name']).categories) #again 811 Stations
print(pd.Categorical(data['Line']).categories) #121 Lines
print(pd.Categorical(data['Platform type']).categories) #5 Types    
print(pd.Categorical(data['Step']).categories) #4 Types
'''