<filename>src/workflow_helpers.py<gh_stars>0
#!/usr/bin/env python3
# -*- coding: utf-8 -*-


import os
from multiprocessing import cpu_count

import numpy as np
import pandas as pd
from IPython.display import display
from joblib import Parallel, delayed

import src.feature_helpers as fh
from src.geopandas_helpers import get_distance


def check_properties(
    df,
    distance_percs_to_show=[
        0.01,
        0.05,
        0.1,
        0.5,
        0.7,
        0.95,
        0.99,
        0.999,
        0.99999,
    ],
    dist_threshold_dict={"99%": 0.8},
    frac_unlisted_stations_thresh=12,
    remove_max_percentile=False,
    quantiles_to_show=[],
):
    """
    Check (1) quantiles for distance between trip departure location and json
    feed location and (2) fraction of trips departing from stations not listed
    in JSON feed
    """
    print("Checking attributes about departure stations...", end="")
    # Distance between trip departure location and json feed location
    distances_listed_stations = df[
        (~df["latitude_feed"].isna()) | (~df["longitude_feed"].isna())
    ][
        [
            "from_station_name",
            "latitude_start",
            "longitude_start",
            "latitude_feed",
            "longitude_feed",
            "start_station_to_feed_km",
        ]
    ]
    if remove_max_percentile:
        max_dist_index = distances_listed_stations[
            "start_station_to_feed_km"
        ].idxmax()
        distances_listed_stations = distances_listed_stations.loc[
            ~distances_listed_stations.index.isin([max_dist_index])
        ]

    distance_percs = distances_listed_stations[
        "start_station_to_feed_km"
    ].describe(percentiles=distance_percs_to_show)

    if quantiles_to_show:
        distance_percs = distance_percs.loc[quantiles_to_show]

    distance_percs = distance_percs.reset_index().rename(
        columns={"index": "statistic"}
    )

    distance_percentile = list(dist_threshold_dict.keys())[0]
    distance_percentile_threshold = list(dist_threshold_dict.values())[0]
    perc_mask = distance_percs["statistic"] == distance_percentile
    percentile = distance_percs.loc[perc_mask, "start_station_to_feed_km"]
    # print(percentile.iloc[0], distance_percentile_threshold)
    try:
        assert percentile.iloc[0] <= distance_percentile_threshold
    except AssertionError as _:
        print(
            "Majority of departure station-to-JSON feed location distances "
            "not within expected percentile threshold"
        )

    # Fraction of trips departing from stations not listed in JSON feed
    no_lat_feed = df["latitude_feed"].isna()
    no_lon_feed = df["longitude_feed"].isna()
    non_missing_station_mask = ~df["from_station_name"].isna()
    trips_from_unlisted_station = df[
        (no_lat_feed) & (no_lon_feed) & (non_missing_station_mask)
    ]
    frac_trips_from_unlisted_stations = 100 * (
        len(trips_from_unlisted_station) / len(df)
    )
    # print(frac_trips_from_unlisted_stations, frac_unlisted_stations_thresh)
    try:
        assert (
            frac_trips_from_unlisted_stations <= frac_unlisted_stations_thresh
        )
    except AssertionError as _:
        print("Fraction of unlisted stations not within expected threshold")
    print("done.")
    return [
        distance_percs,
        trips_from_unlisted_station,
        percentile.iloc[0],
        frac_trips_from_unlisted_stations,
    ]


def featurize(
    data,
    fname,
    f_cntr,
    num_objs,
    max_duration,
    station_renaming_dict,
    non_stations,
    stations_metadata,
    crs,
    epsg,
    use_parallel=True,
):
    if not use_parallel:
        print(
            (
                f"Adding features for dataset {f_cntr+1}/{num_objs} "
                f"({os.path.basename(fname)})...(sizes="
            ),
            end="",
        )
    len_data_1 = len(data)
    if not use_parallel:
        print(f"{len_data_1}, ", end="")
    data, _ = fh.filter_data(
        data.drop_duplicates(ignore_index=True).copy(),
        dur_max=max_duration,
    )
    data = fh.add_datepart(data)
    len_data_2 = len(data)
    if not use_parallel:
        assert len_data_1 >= len_data_2
        print(f"{len_data_2}, ", end="")
    data["from_station_name"] = data["from_station_name"].replace(
        station_renaming_dict
    )
    data = data[~data["from_station_name"].isin(non_stations)]
    len_data_3 = len(data)
    if not use_parallel:
        assert len_data_2 >= len_data_3
        print(f"{len_data_3}, ", end="")
    data = data.merge(
        stations_metadata.rename(
            columns={
                "station_name": "from_station_name",
                "latitude": "latitude_feed",
                "longitude": "longitude_feed",
            }
        )[["from_station_name", "latitude_feed", "longitude_feed"]],
        on="from_station_name",
        how="left",
    )
    len_data_4 = len(data)
    if not use_parallel:
        assert len_data_3 == len_data_4
        print(f"{len_data_4}, ", end="")
    data = get_distance(
        data,
        ["latitude_start", "longitude_start"],
        ["latitude_feed", "longitude_feed"],
        crs=crs,
        epsg=epsg,
    )
    len_data_5 = len(data)
    if not use_parallel:
        assert len_data_3 == len_data_5
        print(f"{len_data_5})...", end="")
        print("done.")
    df_lengths = pd.DataFrame.from_dict(
        {
            "fname": os.path.basename(fname),
            "pre_featurize": len_data_1,
            "post_station_removal": len_data_2,
            "post_duration_filter": len_data_3,
            "post_add_feed_lat_lon": len_data_4,
            "post_add_dist": len_data_5,
        },
        orient="index",
    ).T
    f_cntr += 1
    return [data, df_lengths]


def featurize_2021_trips(
    file_paths,
    objs,
    stations_metadata,
    station_renaming_dict,
    non_stations,
    max_duration=45 * 60,
    crs=4326,
    epsg=6345,
    dtype_changer_dict={},
    use_parallel=True,
):
    """
    Add features to each quarterly/monthly file of bikeshare trips and
    optionally change dtypes (required for saving to parquet file)
    """
    if use_parallel:
        executor = Parallel(n_jobs=cpu_count(), backend="multiprocessing")
        objs_filtered = []
        tasks = (
            delayed(featurize)(
                data,
                fname,
                f_cntr,
                len(objs),
                max_duration,
                station_renaming_dict,
                non_stations,
                stations_metadata,
                crs,
                epsg,
                use_parallel,
            )
            for f_cntr, (fname, data) in enumerate(zip(file_paths, objs))
        )
        objs_filtered = executor(tasks)
    else:
        objs_filtered = [
            featurize(
                data,
                fname,
                f_cntr,
                len(objs),
                max_duration,
                station_renaming_dict,
                non_stations,
                stations_metadata,
                crs,
                epsg,
                use_parallel,
            )
            for f_cntr, (fname, data) in enumerate(zip(file_paths, objs))
        ]
    print("Combining quarterly/monthly datasets...", end="")
    data, data_lengths = [
        pd.concat([obj[k] for obj in objs_filtered], axis=0, ignore_index=True)
        for k in [0, 1]
    ]
    print("done.")

    if dtype_changer_dict:
        print(
            (
                "Changing column datatypes for "
                f"{', '.join(list(dtype_changer_dict))}..."
            ),
            end="",
        )
        data = data.astype(dtype_changer_dict)
        print("done.")
    return [data, data_lengths]


def fill_trip_start_lat_lon(df, d_missing_lat_lon):
    """
    Fill missing values in the latitude_start and longitude_start columns.
    """
    # Get trips with missing lat and lon
    df_no_lat_lon = df[
        df[["latitude_start", "longitude_start"]].isna().all(axis=1)
    ]
    display(
        df_no_lat_lon.head().style.set_caption(
            f"Missing lat and lon for {len(df_no_lat_lon)} trips"
        )
    )

    # Compare list of observed and expected stations with missing trip start
    # lat and lon
    observed_stations_no_trip_lat_lon = (
        df_no_lat_lon["from_station_name"].unique().tolist()
    )
    try:
        assert set(list(d_missing_lat_lon)) == set(
            observed_stations_no_trip_lat_lon
        )
    except AssertionError as _:
        print(
            "Found unexpected list of stations with missing trip starting "
            "latitude and longitude"
        )
        df_no_trip_lat_lon = pd.DataFrame(
            zip(
                *[
                    list(set(observed_stations_no_trip_lat_lon)),
                    set(list(d_missing_lat_lon)),
                ]
            ),
            columns=["observed", "expected"],
        )
        display(
            df_no_trip_lat_lon.style.set_caption(
                "Observed and Expected stations with missing trip start "
                "latitude and longitude"
            )
        )
    else:
        df["latitude_start"] = df["latitude_start"].fillna(
            df["from_station_name"].map(
                {k: v[0] for k, v in d_missing_lat_lon.items()}
            )
        )
        df["longitude_start"] = df["longitude_start"].fillna(
            df["from_station_name"].map(
                {k: v[1] for k, v in d_missing_lat_lon.items()}
            )
        )
        print(
            "Expected list of stations with missing trip starting latitude "
            "and longitude. Filled."
        )
    df_no_lat_lon = df[
        df[["latitude_start", "longitude_start"]].isna().all(axis=1)
    ]
    try:
        assert df_no_lat_lon.empty
    except AssertionError as _:
        print("Found trips missing latitude and longitude")
        display(
            df_no_lat_lon.style.set_caption(
                f"Missing lat and lon for {len(df_no_lat_lon)} trips"
            )
        )
    return df


def process_2020_data_condensed(
    fpath,
    columns_2020,
    twenty_twenty_renaming_dict,
):
    """
    Process single quarter/month of 2020 trips data.
    """
    df = (
        pd.read_csv(
            fpath,
            parse_dates=["started_at", "ended_at"],
        )
        .rename(columns=twenty_twenty_renaming_dict)
        .drop(columns=["rideable_type"], axis=1)
    )
    df["tripduration"] = (
        (df["stoptime"] - df["starttime"]).dt.total_seconds()
    ).astype(float)
    df[["gender", "dpcapacity_end", "dpcapacity_start"]] = [
        [np.nan, np.nan, np.nan]
    ] * len(df)
    df = fh.add_datepart_prelims(df)
    df = df[columns_2020 + ["month", "day", "hour"]]
    return df


def load_all_2020_data(
    fpaths,
    twenty_twenty_renaming_dict,
    columns_2020,
):
    """Process 2020 trips data."""
    executor = Parallel(n_jobs=cpu_count(), backend="multiprocessing")
    tasks = (
        delayed(process_2020_data_condensed)(
            fpath,
            columns_2020,
            twenty_twenty_renaming_dict,
        )
        for fpath in fpaths
    )
    dfs = executor(tasks)
    return dfs


def merge(trip, station):
    """
    Merge trip and station information.
    """
    trip = trip[
        [
            "trip_id",
            "usertype",
            "gender",
            "starttime",
            "stoptime",
            "tripduration",
            "from_station_id",
            "from_station_name",
            "to_station_id",
            "to_station_name",
        ]
    ]
    station = station[["id", "latitude", "longitude", "dpcapacity"]]

    trip_start = pd.merge(
        left=trip,
        right=station,
        how="left",
        left_on="from_station_id",
        right_on="id",
    )
    trip_station = pd.merge(
        left=trip_start,
        right=station,
        how="left",
        left_on="to_station_id",
        right_on="id",
        suffixes=("_start", "_end"),
    )

    columns = [
        "trip_id",
        "usertype",
        "gender",
        "starttime",
        "stoptime",
        "tripduration",
        "from_station_id",
        "from_station_name",
        "latitude_start",
        "longitude_start",
        "dpcapacity_start",
        "to_station_id",
        "to_station_name",
        "latitude_end",
        "longitude_end",
        "dpcapacity_end",
    ]
    trip_station = trip_station[columns]
    trip_station = fh.add_datepart_prelims(trip_station)
    return trip_station


def process_2018_2019_data_condensed(
    fpath,
    twenty_eighteen_nineteen_final_cols,
    twenty_eighteen_nineteen_date_cols_dict,
    df_stations,
    station_renaming_dict,
    non_stations,
):
    """
    Process single quarter of 2018-2019 trips data.
    """
    df = pd.read_csv(
        fpath,
        parse_dates=twenty_eighteen_nineteen_date_cols_dict[
            os.path.basename(fpath)
        ],
    )
    df.columns = twenty_eighteen_nineteen_final_cols
    # print(list(df))
    df["tripduration"] = df["tripduration"].replace(",", "", regex=True)
    df["tripduration"] = df["tripduration"].astype(float)
    df["from_station_name"] = df["from_station_name"].replace(
        station_renaming_dict
    )
    df = df[~df["from_station_name"].isin(non_stations)]
    df_merged = merge(df, df_stations)
    return df_merged


def load_all_2018_2019_data(
    fpaths,
    twenty_eighteen_nineteen_final_cols,
    twenty_eighteen_nineteen_date_cols_dict,
    df_stations,
    station_renaming_dict,
    non_stations,
):
    executor = Parallel(n_jobs=cpu_count(), backend="multiprocessing")
    tasks = (
        delayed(process_2018_2019_data_condensed)(
            fpath,
            twenty_eighteen_nineteen_final_cols,
            twenty_eighteen_nineteen_date_cols_dict,
            df_stations,
            station_renaming_dict,
            non_stations,
        )
        for fpath in fpaths
    )
    dfs = executor(tasks)
    return dfs
