forerad

hourly-volume-rollup (5091B)

---

 #!/usr/bin/env python3
 """
 This script handles the population of the daily_volume_rollup table
 """
 import argparse
 import datetime
 import pandas as pd
 import geopandas as gpd
 import pathlib
 import os
 import multiprocessing
 
 import forerad.scrapers.historical as historical
 import forerad.persistence as persistence
 import forerad.utils as utils
 
 store = persistence.SQLiteStore()
 
 # The day that Citibike began publishing data on
 ORIGIN_DATE = datetime.date(2013, 6, 1)
 REPO_ROOT = pathlib.Path(__file__).parent.parent
 
 def detect_status(year: int, month: int) -> tuple[set[pd.Timestamp], set[pd.Timestamp], set[pd.Timestamp]]:
     """
     Given a month, look at the hourly rollup and generate a set of missing
     members
     """
     first_day = datetime.date(year, month, 1)
     next_month = utils.next_month(first_day.year, first_day.month)
     hours_in_month = pd.date_range(
         start=first_day,
         end=next_month,
         freq="1H",
         tz=utils.TZ_NYC,
         inclusive="left"
     )
     expected = set(hours_in_month.tz_convert(utils.TZ_UTC))
 
     rollup = store.fetch_hourly_volume_rollup(first_day, next_month)
     rollup.index = rollup.index.tz_convert(utils.TZ_UTC)
     actual = set(rollup.index)
 
     missing = expected - actual
 
     return expected, actual, missing
 
 
 def is_complete(archive: historical.HistoricalTripArchive) -> bool:
     _, _, missing = detect_status(archive.date.year, archive.date.month)
     return len(missing) == 0
 
 def derive_rollup(archive: historical.HistoricalTripArchive):
     df = archive.fetch_df()
 
     boroughs = gpd.read_file(REPO_ROOT / 'resources' / 'borough-boundaries.geojson')
     df = gpd.GeoDataFrame(df, geometry=gpd.points_from_xy(df.start_lng, df.start_lat))\
         .set_crs('EPSG:4326')
     df = df.sjoin(boroughs, predicate='within')
 
     df['trip_count'] = 1
     grouper = pd.Grouper(key='started_at', freq='1H')
     df = df.groupby([grouper, 'boro_name']).count().reset_index()
     df = df[['started_at', 'boro_name', 'trip_count']]
     df = df.pivot_table(
         columns=['boro_name'],
         index=['started_at'],
         fill_value=0,
     )
     df.columns = [f'{boro.lower()}_trips' for _, boro in df.columns]
 
     boros = [
         'manhattan', 'bronx', 'queens', 'brooklyn', 'staten_island'
     ]
     for boro in boros:
         # Pour one out for the one borough that this if statement will always
         # hit.
         key = f'{boro}_trips'
         if key not in df.columns:
             df[key] = 0
 
     # Join with a generated range of timestamps to ensure there are no gaps in
     # the resulting rollup
     hours_in_month = pd.DataFrame(
         index=pd.date_range(
             start=df.index.values.min(),
             end=df.index.values.max(),
             freq="1H",
             tz=utils.TZ_UTC,
         ),
     )
     hours_in_month.index.name = 'started_at'
     df = hours_in_month\
         .join(df)\
         .fillna(0)\
         .reset_index()
 
     df['local_date'] = df['started_at']\
         .dt.tz_convert('America/New_York')\
         .dt.strftime('%Y-%m-%d')
 
     return [archive, df]
 
 def main__populate(month_str: str, cores: int):
     archives = historical.HistoricalTripArchive.list_cached()
 
     if month_str is not None:
         year, month = utils.parse_month_str(month_str)
         archives = [a for a in archives if a.date.year == year and a.date.month == month]
 
     if not month_str:
         # Filter out completed rollups
         archives = [a for a in archives if not is_complete(a)]
 
     utils.logger.info(f'Rolling up {len(archives)} months of data')
     with multiprocessing.Pool(cores) as pool:
         if cores == 1:
             iterator = [derive_rollup(a) for a in archives]
         else:
             iterator = pool.imap(derive_rollup, archives)
 
         for archive, df in iterator:
             store.write_hourly_volume_rollup(df)
             utils.logger.info(f"Wrote {len(df)} members to table for {archive.month_str}")
 
 
 def main__status():
     available_archives = historical.HistoricalTripArchive.list_cached()
 
     print('Archive status:')
     for archive in available_archives:
         expected, actual, missing = detect_status(archive.date.year, archive.date.month)
         if len(missing) == 0:
             presence = '[x]'
         elif len(actual) > len(expected):
             presence = '[!]'
         else:
             presence = '[ ]'
 
         print(f'\t{presence} {archive.month_str} ({len(actual)}/{len(expected)})')
 
 
 if __name__ == '__main__':
     parser = argparse.ArgumentParser(
         description="Backfill historical trip data into your data store",
     )
     parser.add_argument('action', choices=['status', 'populate'])
     parser.add_argument('--month', help="The month to populate, in YYYY-MM format")
     parser.add_argument(
         '--cores',
         type=int,
         default=1,
         help="how many cores to process with (be careful with memory requirements!)"
     )
 
     args = parser.parse_args()
 
     if args.action == 'populate':
         main__populate(args.month, args.cores)
     elif args.action == 'status':
         main__status()