Add options to handle larger dataset for location models

activitysim/estimation/larch/location_choice.py

@@ -3,6 +3,8 @@
 import os
 from pathlib import Path
 from typing import Collection
+import pickle
+from datetime import datetime
 
 import numpy as np
 import pandas as pd
@@ -46,6 +48,8 @@ def location_choice_model(
  settings_file="{name}_model_settings.yaml",
  landuse_file="{name}_landuse.csv",
  return_data=False,
+ alt_values_to_feather=False,
+ chunking_size=None,
 ):
  model_selector = name.replace("_location", "")
  model_selector = model_selector.replace("_destination", "")
@@ -59,12 +63,42 @@ def _read_csv(filename, **kwargs):
  filename = filename.format(name=name)
  return pd.read_csv(os.path.join(edb_directory, filename), **kwargs)
 
+ def _read_feather(filename, **kwargs):
+ filename = filename.format(name=name)
+ return pd.read_feather(os.path.join(edb_directory, filename), **kwargs)
+
+ def _to_feather(df, filename, **kwargs):
+ filename = filename.format(name=name)
+ return df.to_feather(os.path.join(edb_directory, filename), **kwargs)
+
+ def _read_pickle(filename, **kwargs):
+ filename = filename.format(name=name)
+ return pd.read_pickle(os.path.join(edb_directory, filename))
+
+ def _to_pickle(df, filename, **kwargs):
+ filename = filename.format(name=name)
+ return df.to_pickle(os.path.join(edb_directory, filename))
+
+ def _file_exists(filename):
+ filename = filename.format(name=name)
+ return os.path.exists(os.path.join(edb_directory, filename))
+
  coefficients = _read_csv(
  coefficients_file,
  index_col="coefficient_name",
  )
  spec = _read_csv(spec_file, comment="#")
- alt_values = _read_csv(alt_values_file)
+
+ # read alternative values either as csv or feather file
+ alt_values_fea_file = alt_values_file.replace(".csv", ".fea")
+ if os.path.exists(
+ os.path.join(edb_directory, alt_values_fea_file.format(name=name))
+ ):
+ alt_values = _read_feather(alt_values_fea_file)
+ else:
+ alt_values = _read_csv(alt_values_file)
+ if alt_values_to_feather:
+ _to_feather(df=alt_values, filename=alt_values_fea_file)
  chooser_data = _read_csv(chooser_file)
  landuse = _read_csv(landuse_file, index_col="zone_id")
  master_size_spec = _read_csv(size_spec_file)
@@ -152,7 +186,48 @@ def _read_csv(filename, **kwargs):
 
  chooser_index_name = chooser_data.columns[0]
  x_co = chooser_data.set_index(chooser_index_name)
- x_ca = cv_to_ca(alt_values.set_index([chooser_index_name, alt_values.columns[1]]))
+
+ def split(a, n):
+ k, m = divmod(len(a), n)
+ return (a[i * k + min(i, m) : (i + 1) * k + min(i + 1, m)] for i in range(n))
+
+ # process x_ca with cv_to_ca with or without chunking
+ x_ca_pickle_file = "{name}_x_ca.pkl"
+ if chunking_size == None:
+ x_ca = cv_to_ca(
+ alt_values.set_index([chooser_index_name, alt_values.columns[1]])
+ )
+ elif _file_exists(x_ca_pickle_file):
+ # if pickle file from previous x_ca processing exist, load it to save time
+ time_start = datetime.now()
+ x_ca = _read_pickle(x_ca_pickle_file)
+ print(
+ f"x_ca data loaded from {name}_x_ca.fea - time elapsed {(datetime.now() - time_start).total_seconds()}"
+ )
+ else:
+ time_start = datetime.now()
+ # calculate num_chunks based on chunking_size (or max number of rows per chunk)
+ num_chunks = int(len(alt_values) / chunking_size)
+ all_person_ids = list(alt_values["person_id"].unique())
+ split_ids = list(split(all_person_ids, num_chunks))
+ x_ca_list = []
+ i = 0
+ for chunk_ids in split_ids:
+ alt_values_i = alt_values[alt_values["person_id"].isin(chunk_ids)]
+ x_ca_i = cv_to_ca(
+ alt_values_i.set_index([chooser_index_name, alt_values_i.columns[1]])
+ )
+ x_ca_list.append(x_ca_i)
+ print(
+ f"\rx_ca_i compute done for chunk {i}/{num_chunks} - time elapsed {(datetime.now() - time_start).total_seconds()}"
+ )
+ i = i + 1
+ x_ca = pd.concat(x_ca_list, axis=0)
+ # save final x_ca result as pickle file to save time for future data loading
+ _to_pickle(df=x_ca, filename=x_ca_pickle_file)
+ print(
+ f"x_ca compute done - time elapsed {(datetime.now() - time_start).total_seconds()}"
+ )
 
  if CHOOSER_SEGMENT_COLUMN_NAME is not None:
  # label segments with names