stop freq

activitysim/estimation/larch/stop_frequency.py

@@ -27,11 +27,29 @@ def stop_frequency_data(
     edb_directory = edb_directory.format(name=name)
 
     settings_file = settings_file.format(name=name)
-    with open(os.path.join(edb_directory, settings_file), "r") as yf:
-        settings = yaml.load(
-            yf,
-            Loader=yaml.SafeLoader,
-        )
+    try:
+        with open(os.path.join(edb_directory, settings_file), "r") as yf:
+            settings = yaml.load(
+                yf,
+                Loader=yaml.SafeLoader,
+            )
+    except FileNotFoundError:
+        # search in all first level subdirectories for the settings file
+        for subdir in os.listdir(edb_directory):
+            if os.path.isdir(os.path.join(edb_directory, subdir)):
+                try:
+                    with open(
+                        os.path.join(edb_directory, subdir, settings_file), "r"
+                    ) as yf:
+                        settings = yaml.load(
+                            yf,
+                            Loader=yaml.SafeLoader,
+                        )
+                        break
+                except FileNotFoundError:
+                    pass
+        else:
+            raise
 
     segments = [i["primary_purpose"] for i in settings["SPEC_SEGMENTS"]]
 
@@ -111,10 +129,30 @@ def stop_frequency_data(
         seg_alt_names_to_codes.append(alt_names_to_codes)
         seg_alt_codes_to_names.append(alt_codes_to_names)
 
-        chooser_data = pd.read_csv(
-            seg_subdir / chooser_data_file.format(name=name),
-            index_col=values_index_col,
-        )
+        # load parquet if available, otherwise pickle, or if all else fails csv
+        if (
+            (seg_subdir / chooser_data_file.format(name=name))
+            .with_suffix(".parquet")
+            .exists()
+        ):
+            chooser_data = pd.read_parquet(
+                (seg_subdir / chooser_data_file.format(name=name)).with_suffix(
+                    ".parquet"
+                ),
+            ).set_index(values_index_col)
+        elif (
+            (seg_subdir / chooser_data_file.format(name=name))
+            .with_suffix(".pkl")
+            .exists()
+        ):
+            chooser_data = pd.read_pickle(
+                (seg_subdir / chooser_data_file.format(name=name)).with_suffix(".pkl"),
+            ).set_index(values_index_col)
+        else:
+            chooser_data = pd.read_csv(
+                seg_subdir / chooser_data_file.format(name=name),
+                index_col=values_index_col,
+            )
         seg_chooser_data.append(chooser_data)
 
     return Dict(