attach skim data

activitysim/workflows/example_runner.yaml

@@ -197,14 +197,19 @@ reporting:
   in:
     title: Trip Distance
 
+- name: activitysim.workflows.steps.contrast.attach_skim_data
+  in:
+    tablename: trips
+    otaz_col: origin
+    dtaz_col: destination
+    time_col: depart
+    skim_vars: DIST
+
 - name: activitysim.workflows.steps.contrast.trip_distance
   in:
     title: Trip Distance by Primary Purpose, <10 miles
     title_level: 3
     grouping: primary_purpose
-    otaz_col: origin
-    dtaz_col: destination
-    time_col: depart
     dist_bins: 20
     dist_skim_name: DIST
     max_dist: 10
@@ -214,9 +219,6 @@ reporting:
     title: Trip Distance by Primary Purpose
     title_level: 3
     grouping: primary_purpose
-    otaz_col: origin
-    dtaz_col: destination
-    time_col: depart
     dist_bins: 20
     dist_skim_name: DIST
 

activitysim/workflows/steps/contrast/attach_skim_data.py

@@ -0,0 +1,78 @@
+import logging
+import pandas as pd
+from pypyr.context import Context
+from ..progression import reset_progress_step
+
+logger = logging.getLogger(__name__)
+
+
+def run_step(context: Context) -> None:
+
+    contrast_data = context.get('contrast_data')
+    skims = context.get('skims')
+    skim_vars = context.get_formatted('skim_vars')
+    tablename = context.get_formatted('tablename')
+    otaz_col = context.get_formatted('otaz_col')
+    dtaz_col = context.get_formatted('dtaz_col')
+    time_col = context.get_formatted('time_col')
+
+    if isinstance(skim_vars, str):
+        skim_vars = [skim_vars]
+    if len(skim_vars) == 1:
+        skim_vars_note = skim_vars[0]
+    else:
+        skim_vars_note = f"{len(skim_vars)} skim vars"
+
+    reset_progress_step(description=f"attach skim data / {tablename} <- {skim_vars_note}")
+
+    contrast_data = attach_skim_data(
+        contrast_data,
+        skims,
+        skim_vars,
+        tablename,
+        otaz_col,
+        dtaz_col,
+        time_col,
+    )
+    context['contrast_data'] = contrast_data
+
+
+
+
+def attach_skim_data(
+    tablesets,
+    skims,
+    skim_vars,
+    tablename,
+    otaz_col,
+    dtaz_col,
+    time_col,
+):
+    if not isinstance(skims, dict):
+        skims = {i: skims for i in tablesets.keys()}
+
+    for key, tableset in tablesets.items():
+        skim_subset = skims[key][skim_vars]
+
+        zone_ids = tableset['land_use'].index
+        if zone_ids.is_monotonic_increasing and zone_ids[-1] == len(zone_ids) + zone_ids[0] - 1:
+            offset = zone_ids[0]
+            looks = [
+                tableset[tablename][otaz_col].rename('otaz') - offset,
+                tableset[tablename][dtaz_col].rename('dtaz') - offset,
+            ]
+        else:
+            remapper = dict(zip(zone_ids, pd.RangeIndex(len(zone_ids))))
+            looks = [
+                tableset[tablename][otaz_col].rename('otaz').apply(remapper.get),
+                tableset[tablename][dtaz_col].rename('dtaz').apply(remapper.get),
+            ]
+        if 'time_period' in skim_subset.dims:
+            looks.append(
+                tableset[tablename][time_col].apply(skims[key].attrs['time_period_imap'].get).rename('time_period'),
+            )
+        look = pd.concat(looks, axis=1)
+        out = skim_subset.iat.df(look)
+        tablesets[key][tablename] = tablesets[key][tablename].assign(**out)
+
+    return tablesets

activitysim/workflows/steps/contrast/trip_distance.py

@@ -20,23 +20,17 @@ def run_step(context: Context) -> None:
     title = context.get_formatted('title') or "Trip Length Distribution"
     title_level = context.get('title_level', None)
 
-    otaz_col = context.get_formatted('otaz_col')
-    dtaz_col = context.get_formatted('dtaz_col')
-    time_col = context.get_formatted('time_col')
     dist_bins = context.get_formatted('dist_bins')
     max_dist = context.get_formatted_or_default('max_dist', None)
 
-    reset_progress_step(description=f"report trip mode choice / {grouping}")
+    reset_progress_step(description=f"report trip distance / {grouping}")
 
     with report:
         report << fig(title, level=title_level)
         report << compare_trip_distance(
             contrast_data,
             skims,
             dist_skim_name,
-            otaz_col=otaz_col,
-            dtaz_col=dtaz_col,
-            time_col=time_col,
             dist_bins=dist_bins,
             grouping=grouping,
             title=None,
@@ -51,9 +45,6 @@ def compare_trip_distance(
     tablesets,
     skims,
     dist_skim_name,
-    otaz_col='origin',
-    dtaz_col='destination',
-    time_col='depart',
     dist_bins=20,
     grouping='primary_purpose',
     title="Trip Length Distribution",
@@ -65,33 +56,13 @@ def compare_trip_distance(
 
     distances = {}
     for key, tableset in tablesets.items():
-        skim_dist = skims[key][[dist_skim_name]]
-
-        zone_ids = tableset['land_use'].index
-        if zone_ids.is_monotonic_increasing and zone_ids[-1] == len(zone_ids) + zone_ids[0] - 1:
-            offset = zone_ids[0]
-            looks = [
-                tableset['trips'][otaz_col].rename('otaz') - offset,
-                tableset['trips'][dtaz_col].rename('dtaz') - offset,
-            ]
-        else:
-            remapper = dict(zip(zone_ids, pd.RangeIndex(len(zone_ids))))
-            looks = [
-                tableset['trips'][otaz_col].rename('otaz').apply(remapper.get),
-                tableset['trips'][dtaz_col].rename('dtaz').apply(remapper.get),
-            ]
-        if 'time_period' in skim_dist.dims:
-            looks.append(
-                tableset['trips'][time_col].apply(skims[key].attrs['time_period_imap'].get).rename('time_period'),
-            )
-        look = pd.concat(looks, axis=1)
-        distances[key] = skims[key][[dist_skim_name]].iat.df(look)
+        distances[key] = tableset['trips'][dist_skim_name]
 
     if dist_bins is not None:
         result = pd.concat(distances, names=['source'])
         if max_dist is not None:
             result = result[result <= max_dist]
-        result = pd.cut(result.iloc[:, 0], dist_bins).to_frame()
+        result = pd.cut(result, dist_bins).to_frame()
         distances = {k:result.loc[k] for k in tablesets.keys()}
 
     data = {}

activitysim/workflows/steps/progression.py

@@ -82,8 +82,10 @@ def update_progress_overall(*args, **kwargs):
 
 
 def reset_progress_step(*args, description='', prefix='', **kwargs):
-    print("╭" + "─" * (len(description)+2) + "╮")
-    print(f"│ {description} │ {time.strftime('%I:%M:%S %p')}")
-    print("╰" + "─" * (len(description)+2) + "╯")
     if not os.environ.get('NO_RICH', False):
+        print(f"\u23F1  {time.strftime('%I:%M:%S %p')} - {description}")
         progress.reset(progress_step, *args, description=prefix+description, **kwargs)
+    else:
+        print("╭" + "─" * (len(description) + 2) + "╮")
+        print(f"│ {description} │ {time.strftime('%I:%M:%S %p')}")
+        print("╰" + "─" * (len(description) + 2) + "╯")