"""Kedro pipeline nodes for the ``electrify_trips`` pipeline (Model Module 4).
This module implements the per-vehicle charging-choice simulation described in
**Module 4 ("Simulate electrified dwells")** of Passow & Rajagopal (2026).
Starting from a dataset of observed diesel-HDT dwell events, it produces a
dataset of *electrified* dwells annotated with charging decisions — which mode
was used, how much energy was transferred, and how much delay was incurred —
for a single scenario (state of the world).
Pipeline overview
-----------------
The nodes in this module are executed in the following logical order by the
Kedro pipeline:
1. **filter_vehicles** — restrict the dwell data to the vehicle cohort
selected for the current scenario.
2. **calc_vehicle_ranges** — assign each vehicle a design range (miles) and
battery capacity (kWh) derived from its observed shift mileage
distribution.
3. **calc_dwell_durations** — shrink each dwell window by plug-in/plug-out
overhead and compute net available charging time.
4. **prepare_modes** / **assign_modes** / **build_mode_power_lut** — annotate
each dwell with the charging modes available at that location and the
maximum deliverable power.
5. **merge_dwellset_node** — join per-vehicle parameters (e.g. consumption
rate) into the dwell table.
6. **calc_energy_use** — compute per-trip energy demand (kWh).
7. **mark_critical_days** — classify vehicle-days where en-route charging is
necessary because total shift energy exceeds battery capacity.
8. **filter_dwells** — drop en-route dwells that are too short or on
non-critical days (pre-simulation filter).
9. **mark_shift_powers** — record the maximum power available later in the
shift at each dwell, used as a look-ahead input to the charging algorithm.
10. **simulate_charging_choice** — run the forward-looking utility-
maximisation charging-choice algorithm (Numba JIT-compiled) for every
vehicle.
11. **filter_dwells_post** — drop optional stops where the simulation
determined no charging occurred (post-simulation filter).
Key design decisions
--------------------
- **Critical-day filtering**: En-route (truck-stop) charging is only
considered on vehicle-days where the total shift energy demand exceeds
battery capacity. This reduces the search space and avoids over-estimating
public-charging demand on days when depot charging alone suffices.
- **Optional stops**: Proxy dwells inserted at truck stops along routed paths
have zero net duration. They are carried through the pre-simulation filter
and dropped post-simulation if the vehicle chose not to charge there.
- **Dask support**: All nodes handle both pandas- and Dask-backed
``DwellSet`` objects. Dask execution partitions the vehicle fleet so that
each partition is processed independently; callers must ensure vehicle
sequences are not split across partitions.
- **Numba JIT compilation**: The inner charging loop in
``simulate_charging_choice`` is JIT-compiled. An optional pre-compilation
step warms up the compiled functions before the main run.
References
----------
Passow, F., & Rajagopal, R. (2026). Identifying indicators to inform
proactive substation upgrades for charging electric heavy-duty trucks.
*Applied Energy* (submitted March 2026).
Liu, J., et al. Utility-maximisation charging choice model (inspiration for
``ForwardLookingChargingChoiceStrategy``).
"""
from __future__ import annotations
import logging
import dask.dataframe as dd
import geopandas as gpd
import numpy as np
import pandas as pd
from laurel.models.charging_algorithms import ForwardLookingChargingChoiceStrategy
from laurel.models.dwell_sets import CumAggFunc, DwellSet
from laurel.utils.data import merge_dataframes_node
from laurel.utils.mode_masks import bool_arr_to_bits
from laurel.utils.time import total_hours
logger = logging.getLogger(__name__)
[docs]
def drop_cols_to_pandas(
vehs: pd.DataFrame | gpd.GeoDataFrame, params: dict
) -> pd.DataFrame:
"""Drop columns unneeded downstream and coerce the result to a plain pandas DataFrame.
Geometry-bearing GeoDataFrames (e.g. from ``vehicles_labelled``) carry columns
that Dask cannot pyarrow-string-encode (raw WKB/shapely objects). This node drops
those columns and returns a plain DataFrame so that later Dask operations (e.g.
``dd.from_pandas`` in ``simulate_charging_choice``) never see them.
Args:
vehs: DataFrame or GeoDataFrame to process.
params: Configuration dict with the following key:
- ``drop_cols`` (list[str]): Column names to drop.
Returns:
Plain ``pd.DataFrame`` with ``drop_cols`` removed.
"""
vehs = vehs.drop(columns=params["drop_cols"])
return pd.DataFrame(vehs)
[docs]
def filter_vehicles(dw: DwellSet, vehs: pd.DataFrame, params: dict) -> DwellSet:
"""Filter the dwell data to only include vehicles present in the vehicles table.
Retains only dwell rows whose vehicle index appears in ``vehs``. Any
vehicle in ``vehs`` that has no dwell rows generates a warning. For
Dask-backed DwellSets the filtered result is repartitioned to a fixed
number of partitions to rebalance work after dropping rows.
Args:
dw: The dwell dataset to filter. Modified in-place and returned.
vehs: Vehicle-level table whose index contains the set of vehicle IDs
to keep. Rows in ``dw`` whose vehicle ID is absent from this index
are dropped.
params: Pipeline parameters. Expected keys:
- ``n_partitions`` (int): Target partition count used when
repartitioning a Dask-backed ``dw`` after filtering.
Returns:
The filtered ``DwellSet`` containing only dwells for vehicles found in
``vehs``.
"""
logger.info("Filter by vehicles by direct dropping")
if not dw.is_dask:
old_len = len(dw.data)
if not dw.is_dask:
keep_idx = np.intersect1d(dw.data.index.values, vehs.index.values)
dw.data = dw.data.loc[keep_idx]
else:
veh_idx = vehs.index.to_frame(index=False)
dw.data = dw.data.merge(
veh_idx, left_index=True, right_on=vehs.index.name, how="inner"
)
dw.data = dw.data.drop(columns=[vehs.index.name])
if not dw.is_dask:
num_no_dwell_vehs = np.setdiff1d(vehs.index.values, keep_idx).size
if num_no_dwell_vehs > 0:
logger.warning(
f"{num_no_dwell_vehs} vehicles were not found in the dwell data."
)
new_len = len(dw.data)
abs_diff = old_len - new_len
pct_diff = round(abs_diff / old_len * 100, 1)
logger.info(f"Rows dropped: {abs_diff}, {pct_diff}%")
if dw.is_dask:
dw.data = dw.data.repartition(npartitions=params["n_partitions"])
return dw
[docs]
def calc_vehicle_ranges(vehs: pd.DataFrame, dw: DwellSet, params: dict) -> pd.DataFrame:
"""Assign a design range (miles) and battery capacity (kWh) to each vehicle.
Design range is determined by taking the maximum of two per-vehicle
criteria computed from observed shift mileage distributions:
1. **Death range** — the ``no_death_shift_frac`` quantile of each
vehicle's single-shift longest-trip distances. The vehicle must be
able to complete any shift's longest leg without running out of charge.
2. **Charge range** — the ``no_charge_shift_frac`` quantile of each
vehicle's total-shift mileage, divided by the usable SoC band
(``soc_buffer_high`` − ``soc_buffer_low``). The vehicle should be
able to complete a typical full shift on one charge.
The continuous desired range is then rounded *up* to the nearest value in
``range_options_miles`` (via ``pd.cut`` with the top bin open), and
multiplied by the vehicle's energy-consumption rate to obtain battery
capacity.
Args:
vehs: Vehicle-level table to augment. Two columns are added in-place:
the design range column and the battery capacity column (names
taken from ``params``). The vehicle index is used to join shift
statistics back to vehicles.
dw: Dwell dataset containing trip-distance and shift-ID columns used
to derive per-vehicle shift mileage statistics.
params: Pipeline parameters. Expected keys:
- ``columns`` (dict): Column-name mappings with sub-keys:
- ``shift`` — shift identifier column in ``dw.data``
- ``range_mi`` — output design-range column name in ``vehs``
- ``batt_kwh`` — output battery-capacity column name in ``vehs``
- ``consump_kwh_per_mi`` — energy-consumption-rate column in
``vehs``
- ``no_death_shift_frac`` (float): Quantile (0–1) of the
longest-trip distribution used for the death-range criterion.
- ``no_charge_shift_frac`` (float): Quantile (0–1) of the
total-shift-miles distribution used for the charge-range
criterion.
- ``soc_buffer_high`` (float): Upper SoC target (fraction, 0–1).
- ``soc_buffer_low`` (float): Lower SoC buffer (fraction, 0–1).
- ``range_options_miles`` (list[float]): Ordered list of
candidate design ranges in miles (e.g. ``[150, 300, 500]``).
The top entry acts as the ceiling; everything above the
second-to-last bin edge is mapped to the top option.
Returns:
The ``vehs`` DataFrame with two new columns appended: design range
(miles) and battery capacity (kWh).
"""
pcols = params["columns"]
# Describe shifts
shift_miles = dw.data.groupby([dw.veh, pcols["shift"]])[dw.trip_dist].agg(
["sum", "max"]
)
renamer = {"sum": "shift_total_miles", "max": "shift_longest_trip"}
shift_miles = shift_miles.rename(columns=renamer)
if dw.is_dask:
shift_miles = shift_miles.compute()
shift_miles = shift_miles.reset_index()
# Compute desired ranges for each criterion
veh_grp = shift_miles.groupby(dw.veh)
kws = {"interpolation": "linear"}
q_dth = params["no_death_shift_frac"]
death_range = veh_grp["shift_longest_trip"].quantile(q=q_dth, **kws)
q_chg = params["no_charge_shift_frac"]
charge_range = veh_grp["shift_total_miles"].quantile(q=q_chg, **kws)
desired_soc_band = params["soc_buffer_high"] - params["soc_buffer_low"]
charge_range = charge_range / desired_soc_band
# Compute design range from desired ranges for each criterion
ranges = pd.concat([death_range, charge_range], axis=1)
ranges["range_desired"] = ranges.max(axis=1)
# Omitting top bin to automatically do the two-way rounding
ropts = params["range_options_miles"]
range_bins = [0.0] + ropts[:-1] + [np.inf]
veh_bins = pd.cut(ranges["range_desired"], bins=range_bins, labels=ropts)
# Assign ranges to vehicles
vehs[pcols["range_mi"]] = vehs.index.map(veh_bins).astype(float)
# Assign battery capacities to vehicles matching ranges
vehs[pcols["batt_kwh"]] = (
vehs[pcols["range_mi"]] * vehs[pcols["consump_kwh_per_mi"]]
)
return vehs
[docs]
def calc_dwell_durations(dw: DwellSet, params: dict) -> DwellSet:
"""Compute net charging-available dwell duration for each stop.
Converts raw plug-in and plug-out overhead times from a numeric unit to
``timedelta``, then shrinks each dwell's start and end times inward by
those overhead amounts. The remaining window represents the time actually
available for charging. Optional stops — identified by identical start
and end times — are left with zero duration so they are not double-counted.
The resulting net dwell duration (in hours, as a float) is written to a
new column.
Args:
dw: Dwell dataset. ``dw.start`` and ``dw.end`` are timestamp columns
that bound each dwell. Modified in-place and returned.
params: Pipeline parameters. Expected keys:
- ``in_out_time_cols`` (dict): Maps ``"plug_in"`` and
``"plug_out"`` to the column names in ``dw.data`` that hold the
overhead durations (in the unit given by ``in_out_time_unit``).
- ``in_out_time_unit`` (str): Time unit string recognised by
``pd.to_timedelta`` / ``dd.to_timedelta`` (e.g. ``"s"`` for
seconds, ``"min"`` for minutes).
- ``dwell_time_col`` (str): Name of the output column that will
hold the net dwell duration in hours.
Returns:
The updated ``DwellSet`` with adjusted ``dw.start`` / ``dw.end``
timestamps and a new float column for net dwell hours.
"""
iocols = params["in_out_time_cols"]
iounit = params["in_out_time_unit"]
for col in iocols.values():
if dw.is_dask:
dw.data[col] = dd.to_timedelta(dw.data[col], unit=iounit)
else:
dw.data[col] = pd.to_timedelta(dw.data[col], unit=iounit)
# Adjust dwell start and end times to allow time for vehicle to plug in and out
# If the stop is optional (proxied by identical start and end times), then leave
# its duration as zero.
is_optional = dw.data[dw.start] == dw.data[dw.end]
dw.data[dw.end] -= ~is_optional * dw.data[iocols["plug_out"]]
dw.data[dw.start] += ~is_optional * dw.data[iocols["plug_in"]]
dwell_time_col = params["dwell_time_col"]
dw.data[dwell_time_col] = total_hours(dw.data[dw.end] - dw.data[dw.start])
return dw
[docs]
def prepare_modes(modes: dict) -> pd.DataFrame:
"""Convert the charging-modes parameter dict into a tidy DataFrame.
The ``modes`` parameter dict maps mode names to their attribute dicts
(e.g. maximum power). Two special keys — ``name_column`` and
``id_column`` — control the column and index names of the resulting table
and are removed before conversion.
Example input (YAML)::
modes:
name_column: mode_name
id_column: mode_id
depot:
max_power_kw: 150
truck_stop:
max_power_kw: 350
Produces a DataFrame with index name ``mode_id``, a ``mode_name`` column
(``"depot"``, ``"truck_stop"``, …), and one column per attribute.
Args:
modes: Charging-modes parameter dictionary, typically loaded from the
Kedro parameters YAML. Must contain the two special keys
``name_column`` and ``id_column``; all remaining keys are treated
as mode names whose values are attribute dicts.
Returns:
A ``pd.DataFrame`` with one row per charging mode, indexed by a
sequential integer ID, and columns for the mode name and each
mode attribute.
"""
modes_copy = dict(modes) # avoid mutating input params dict
name_col = modes_copy.pop("name_column")
id_col = modes_copy.pop("id_column")
modes_df = pd.DataFrame.from_dict(data=modes_copy, orient="index")
modes_df.index.name = name_col
modes_df = modes_df.reset_index()
modes_df.index.name = id_col
return modes_df
[docs]
def assign_modes(dw: DwellSet, modes: pd.DataFrame, params: dict) -> DwellSet: # noqa: PLR0912
"""Annotate each dwell with available charging modes and maximum deliverable power.
Charging-mode availability at a dwell is determined by two independent
mechanisms and then encoded into a compact integer bitmask:
1. **Location-based modes** (e.g. truck-stop charging): A boolean column
is created for each mode whose availability depends on the dwell's
location group. The selector can be a literal ``True``/``False`` or a
dict specifying which ``loc_groups`` values enable the mode (with an
optional ``invert_selection`` flag).
2. **Vehicle-based depot mode**: A single mode whose availability depends
on whether the vehicle's home-base ratio exceeds a threshold. The
ratio column is supplied via ``params["veh_based_mode_avail"]``.
After building one boolean column per mode, all boolean columns are
combined into an integer bitmask (via ``bool_arr_to_bits``), and the
maximum power deliverable under each possible bitmask combination is
looked up from a pre-built LUT (see ``build_mode_power_lut``). The
individual boolean mode columns are then dropped.
Args:
dw: Dwell dataset to annotate. Modified in-place and returned.
modes: Modes table produced by ``prepare_modes``, with one row per
charging mode and columns for mode name and maximum power.
params: Pipeline parameters. Expected keys:
- ``mode_col`` (str): Column in ``modes`` holding mode names.
- ``loc_based_mode_avail`` (dict): Maps mode name → selector.
Each selector is either a ``bool`` (applies globally) or a dict
with ``loc_groups`` (list of location-group values that enable
the mode) and optionally ``invert_selection`` (bool).
- ``loc_group_col`` (str): Column in ``dw.data`` containing the
location-group identifier for each dwell.
- ``veh_based_mode_avail`` (dict): Configuration for the one
vehicle-home-based mode, with sub-keys ``mode_name`` (str),
``ratio_col`` (str), and ``ratio_thresh`` (float).
- ``mode_mask_col`` (str): Output column name for the integer
mode-availability bitmask.
- ``max_power_source_col`` (str): Column in ``modes`` holding
each mode's maximum power (kW).
- ``max_power_col`` (str): Output column name for the maximum
deliverable power at each dwell.
Returns:
The annotated ``DwellSet`` with a mode-availability bitmask column
and a maximum-power column added; individual per-mode boolean columns
are removed.
Raises:
ValueError: If any mode name in ``modes`` already exists as a column
in ``dw.data``, if a mode specified in ``loc_based_mode_avail``
is not present in ``modes``, if a selector value is neither
``bool`` nor ``dict``, or if any mode remains unassigned after
processing both location- and vehicle-based rules.
"""
all_modes = modes[params["mode_col"]].to_list()
modes_in_use = [mode for mode in all_modes if mode in dw.data.columns]
if len(modes_in_use) > 0:
raise ValueError(f"Mode columns already in use: {modes_in_use}")
# Build one boolean column for each location-only mode
mode_locs = params["loc_based_mode_avail"]
for mode, selector in mode_locs.items():
if mode not in all_modes:
raise ValueError(f"Charging mode '{mode}' not available.")
if isinstance(selector, bool):
dw.data[mode] = selector
elif isinstance(selector, dict):
locs = selector["loc_groups"]
invert = selector.get("invert_selection", False)
dw.data[mode] = dw.data[params["loc_group_col"]].isin(locs)
if invert:
dw.data[mode] = ~dw.data[mode]
else:
raise ValueError(
"Available modes must be boolean or dict of location groups."
)
# Add a boolean column for the vehicle-based depot mode
mode_vehs = params["veh_based_mode_avail"]
vmode_col = mode_vehs["mode_name"]
if vmode_col not in dw.data.columns:
dw.data[vmode_col] = True
orig_cols = list(dw.data.columns)
dw.data[vmode_col] &= dw.data[mode_vehs["ratio_col"]] > mode_vehs["ratio_thresh"]
if dw.is_dask: # Reset column ordering to keep in sync with the Dask Dataframe meta
dw.data = dw.data.map_partitions(
lambda part: part.loc[:, orig_cols], meta=dw.data._meta
)
# Check that all modes are accounted for
modes_missed = [mode for mode in all_modes if mode not in dw.data.columns]
if len(modes_missed) > 0:
raise ValueError(f"The following charging modes are unassigned: {modes_missed}")
# Build mode bitmask column
if dw.is_dask:
all_mask = bool_arr_to_bits(np.ones(shape=(len(modes),), dtype=np.bool_))
meta = dw.data._meta.assign(mode_mask_bits=all_mask.dtype)
dw.data = dw.data.map_partitions(
lambda part: part.assign(
mode_mask_bits=bool_arr_to_bits(
part[all_modes].to_numpy(dtype=np.bool_)
)
),
meta=meta,
)
else:
dw.data["mode_mask_bits"] = bool_arr_to_bits(
dw.data[all_modes].to_numpy(dtype=np.bool_)
)
dw.data = dw.data.rename(columns={"mode_mask_bits": params["mode_mask_col"]})
# Build maximum power column
power_mapper = build_mode_power_lut(
mode_names=modes[params["mode_col"]],
mode_powers=modes[params["max_power_source_col"]],
)
if dw.is_dask:
meta = (params["max_power_col"], modes[params["max_power_source_col"]].dtype)
dw.data[params["max_power_col"]] = dw.data[params["mode_mask_col"]].map(
power_mapper, meta=meta
)
else:
dw.data[params["max_power_col"]] = dw.data[params["mode_mask_col"]].map(
power_mapper
)
# Drop boolean columns
dw.data = dw.data.drop(columns=all_modes)
return dw
[docs]
def build_mode_power_lut(
mode_names: pd.Series, mode_powers: pd.Series
) -> dict[int, float]:
"""Pre-compute the maximum deliverable power for every mode-availability bitmask.
Enumerates all 2^N combinations of N charging modes (where N is
``len(mode_names)``), encodes each combination as an integer bitmask via
``bool_arr_to_bits``, and records the maximum mode power across all
*enabled* modes in that combination. A bitmask of zero (no modes
available) maps to 0.0 kW.
This lookup table is used by ``assign_modes`` so that per-dwell maximum
power can be obtained by a single O(1) dictionary lookup rather than
recomputing the max over enabled modes for every row.
Args:
mode_names: Ordered series of charging-mode name strings. The order
determines bit positions in the bitmask (first name → least
significant bit).
mode_powers: Series of maximum power values (kW), aligned with
``mode_names`` by position.
Returns:
A ``dict`` mapping each integer bitmask (``int``) to the maximum
deliverable power (``float``, kW) across all modes enabled in that
bitmask.
"""
name_ls = mode_names.tolist()
power_arr = mode_powers.to_numpy(dtype=float)
combos = pd.DataFrame(
np.array(np.meshgrid(*[[False, True]] * len(name_ls))).T.reshape(
-1, len(name_ls)
),
columns=name_ls,
dtype=bool,
)
bitmask_col = "mode_mask_bits"
combos[bitmask_col] = bool_arr_to_bits(combos[name_ls].to_numpy(dtype=np.bool_))
combos["max_power"] = (
combos[name_ls].to_numpy(dtype=float) * power_arr[None, :]
).max(axis=1)
res_ser = pd.Series(
combos["max_power"].to_numpy(dtype=float),
index=combos[bitmask_col].to_numpy(),
dtype=float,
)
res_dict = res_ser.to_dict()
return res_dict
[docs]
def merge_dwellset_node(dw: DwellSet, right: pd.DataFrame, params: dict) -> DwellSet:
"""Merge a DataFrame into a DwellSet using the standard merge-node helper.
A thin wrapper around ``merge_dataframes_node`` that extracts the
underlying DataFrame from ``dw``, performs the merge, and writes the
result back. All merge semantics (join type, key columns, etc.) are
controlled by ``params`` exactly as they would be for a plain DataFrame
merge node.
Args:
dw: Dwell dataset whose ``data`` attribute is used as the left side
of the merge. Modified in-place and returned.
right: DataFrame to merge in on the right side.
params: Merge parameters forwarded verbatim to
``merge_dataframes_node``. Refer to that helper for the full
parameter schema.
Returns:
The updated ``DwellSet`` with ``dw.data`` replaced by the merged
result.
"""
dw.data = merge_dataframes_node(
left=dw.data,
right=right,
params=params,
)
return dw
[docs]
def calc_energy_use(dw: DwellSet, params: dict) -> DwellSet:
"""Compute per-trip energy demand (kWh) as trip distance × consumption rate.
Multiplies the trip-distance column (``dw.trip_dist``) by a
vehicle-specific energy-consumption rate column to produce per-trip energy
demand in kWh. Both columns must already be present in ``dw.data``; the
result is written to a new column.
Args:
dw: Dwell dataset containing trip distances and energy-consumption
rates. Modified in-place and returned.
params: Pipeline parameters. Expected keys:
- ``energy_col`` (str): Name for the output energy column (kWh).
- ``consump_col`` (str): Name of the column in ``dw.data`` holding
each vehicle's energy consumption rate (kWh/mile).
Returns:
The updated ``DwellSet`` with the new energy column appended.
"""
dw.data[params["energy_col"]] = (
dw.data[dw.trip_dist] * dw.data[params["consump_col"]]
)
return dw
[docs]
def mark_critical_days(dw: DwellSet, params: dict) -> DwellSet:
"""Classify each dwell as belonging to a critical or non-critical vehicle-day.
A vehicle-day is **critical** when the total energy demand for all
remaining trips in the shift exceeds the vehicle's battery capacity,
making en-route charging necessary. Non-critical days can
be completed on a single depot/destination charge, so en-route dwells on
those days are candidates for filtering.
The algorithm proceeds in three steps:
1. **Identify refresh boundaries** — a refresh dwell is one that has
sufficient dwell duration to fully recharge the battery (duration ≥
battery capacity / max power). These mark the segment boundaries
within which energy demand is accumulated.
2. **Accumulate remaining-shift energy** — starting from each refresh
boundary and working backwards, the energy demand of subsequent trips
is summed up to the next refresh boundary. This gives the energy a
vehicle would need on-board if it arrived at the refresh stop empty.
3. **Classify and propagate** — a dwell is initially critical if its
accumulated remaining-shift energy exceeds battery capacity. The
critical flag is then forward-filled within each vehicle's sequence so
that *all* dwells between a refresh boundary and the critical trip
inherit the flag. Partial days (no prior refresh boundary) are
conservatively treated as critical.
Args:
dw: Dwell dataset with energy, duration, power, battery-capacity, and
refresh columns already populated. Modified in-place and returned.
params: Pipeline parameters. Expected keys:
- ``refresh_col`` (str): Boolean column marking refresh-eligible
dwells (e.g. truck-stop dwells with sufficient dwell time).
- ``crit_bound_col`` (str): Temporary column name used internally
for the refresh-and-can-charge boundary flag.
- ``batt_cap_col`` (str): Column holding each dwell row's vehicle
battery capacity (kWh).
- ``max_power_col`` (str): Column holding maximum available
charging power at each dwell (kW).
- ``dur_col`` (str): Column holding net dwell duration (hours).
- ``energy_col`` (str): Column holding per-trip energy demand
(kWh).
- ``energy_col_next_trip`` (str): Temporary column for the next
trip's energy demand.
- ``energy_col_remain_shift`` (str): Output column for total
remaining-shift energy demand (kWh).
- ``crit_col`` (str): Output boolean column marking critical-day
dwells.
Returns:
The updated ``DwellSet`` with a new boolean ``crit_col`` column and
the intermediate ``crit_bound_col`` column removed.
"""
refr_col = params["refresh_col"]
crit_bnd_col = params["crit_bound_col"]
hrs_to_fill = dw.data[params["batt_cap_col"]] / dw.data[params["max_power_col"]]
can_fully_charge = dw.data[params["dur_col"]] >= hrs_to_fill
dw.data[crit_bnd_col] = dw.data[refr_col] & can_fully_charge
nrg_col_next = params["energy_col_next_trip"]
nrg_col_shift = params["energy_col_remain_shift"]
nrg_col_cur = params["energy_col"]
kws = {}
if dw.is_dask:
kws.update({"meta": ("x", "f8")})
# Using no "fill_value" in shift is okay here because the column is NaN-able (float)
dw.data[nrg_col_next] = dw.data.groupby(dw.veh)[nrg_col_cur].shift(-1, **kws)
dw.data[nrg_col_next] = dw.data.groupby(dw.veh)[nrg_col_next].ffill()
dw.accum_masked(
crit_bnd_col,
accum_cols=nrg_col_next,
reverse=True,
write_all=True,
inplace=True,
)
dw.data = dw.data.rename(columns={f"{nrg_col_next}_{crit_bnd_col}": nrg_col_shift})
# Apply vehicle-specific battery capacity
crcol = params["crit_col"]
dw.data[crcol] = dw.data[nrg_col_shift] > dw.data[params["batt_cap_col"]]
# Assume a critical day for partial days, since the point of the critical days
# assumption is to reduce public charging on days when we're sure it's unnecessary
# Note that "boolean" is different from "bool" data type. See Pandas BooleanArray
is_critical = dw.data[crcol].astype("boolean")
is_crit_bnd = dw.data[crit_bnd_col].astype("boolean")
dw.data[crcol] = (is_crit_bnd & is_critical) ^ ~(is_crit_bnd | pd.NA)
dw.data[crcol] = dw.data[crcol].groupby(dw.veh, sort=False).ffill()
dw.data[crcol] = dw.data[crcol].fillna(True).astype(bool)
# dw.data[refr_col] = dw.data[refr_col].fillna(False).astype(bool)
if dw.is_dask:
dw.data = dw.data.drop(columns=[crit_bnd_col])
else:
dw.data.drop(columns=[crit_bnd_col], inplace=True)
return dw
def _filter_dwells_core(
dw: DwellSet,
keep_col: str,
accum_cols_fw_extra: list[str] | None = None,
accum_cols_rv: list[str] | None = None,
drop_cols_extra: list[str] | None = None,
) -> DwellSet:
"""Drop unwanted dwell rows and propagate corrected cumulative columns.
This is the shared implementation called by both ``filter_dwells`` and
``filter_dwells_post``. It wraps ``DwellSet.accum_masked()`` with the
bookkeeping needed to:
- Accumulate trip-distance, trip-duration, and reset columns (and any
extra columns) through the mask so that the values on *kept* rows
correctly reflect the skipped rows.
- Drop rows where ``keep_col`` is ``False``.
- Rename the accumulated ``<col>_<keep_col>`` output columns back to their
original names so downstream nodes are unaffected.
The base set of forward-accumulated columns is always
``[dw.trip_dist, dw.trip_dur, dw.reset]``. Callers may supply additional
forward (``accum_cols_fw_extra``) and reverse (``accum_cols_rv``)
accumulation columns.
Args:
dw: Dwell dataset to filter. Modified in-place and returned.
keep_col: Name of the boolean column in ``dw.data`` indicating which
rows to keep (``True``) or drop (``False``).
accum_cols_fw_extra: Additional column names to accumulate in the
forward direction (i.e. values propagate from preceding kept rows
to the next kept row). ``None`` means no extras.
accum_cols_rv: Column names to accumulate in the *reverse* direction
(i.e. values propagate from following kept rows back to the
current row). ``None`` means no reverse accumulation.
drop_cols_extra: Additional column names to drop after filtering,
beyond ``keep_col`` and the original (pre-rename) accumulated
columns. ``None`` means no extras.
Returns:
The filtered and updated ``DwellSet`` with corrected cumulative
columns and dropped rows.
"""
accum_cols = [dw.trip_dist, dw.trip_dur, dw.reset]
revs = [False] * len(accum_cols)
if accum_cols_fw_extra is not None:
accum_cols += accum_cols_fw_extra
revs += [False] * len(accum_cols_fw_extra)
if accum_cols_rv is not None:
accum_cols += accum_cols_rv
revs += [True] * len(accum_cols_rv)
dw.accum_masked(keep_col, accum_cols=accum_cols, reverse=revs, inplace=True)
dw.data[keep_col] = dw.data[keep_col].astype("boolean")
dw.data[keep_col] = dw.data[keep_col].replace(False, pd.NA)
if dw.is_dask:
dw.data = dw.data.dropna(subset=keep_col)
else:
dw.data = dw.data.dropna(subset=keep_col)
dw.data[keep_col] = dw.data[keep_col].astype(bool)
drop_cols = [keep_col] + accum_cols
if drop_cols_extra is not None:
drop_cols.extend(drop_cols_extra)
dw.data = dw.data.drop(columns=drop_cols)
renamer = {f"{old}_{keep_col}": old for old in accum_cols}
dw.data = dw.data.rename(columns=renamer)
dw.data[dw.reset] = dw.data[dw.reset].astype(bool)
return dw
[docs]
def filter_dwells(dw: DwellSet, params: dict) -> DwellSet:
"""Drop en-route dwells that cannot meaningfully contribute to charging.
Applied *before* the charging-choice simulation. A dwell is dropped when:
- Its net available duration is negative (shorter than plug-in +
plug-out overhead), OR
- (When ``filter_critical_days`` is enabled) it is neither a refresh
location nor part of a critical day, AND it is not an optional stop.
Keeping optional stops (zero-duration proxy stops inserted along routes)
regardless of the critical-day flag preserves them for the
post-simulation filter, which uses actual charging decisions to decide
whether they were visited.
Dropped rows are handled via ``_filter_dwells_core``, which accumulates
trip-distance and related columns through the mask before removing rows
so that successive kept rows maintain correct running totals.
Args:
dw: Dwell dataset to filter. Modified in-place and returned.
params: Pipeline parameters. Expected keys:
- ``dwell_time_col`` (str): Column holding net dwell duration in
hours (negative means too short to plug in/out).
- ``filter_critical_days`` (bool): Whether to apply the
critical-day filter in addition to the duration filter.
- ``filter_cols`` (dict): Required when ``filter_critical_days``
is ``True``. Sub-keys:
- ``refresh`` — boolean column marking refresh-eligible dwells.
- ``crit`` — boolean column marking critical-day dwells.
- ``accum_cols_forward_extra`` (list[str]): Extra columns to
accumulate in the forward direction through the mask.
- ``accum_cols_reverse`` (list[str]): Columns to accumulate in
the reverse direction through the mask.
- ``drop_cols`` (list[str]): Additional columns to drop after
filtering.
Returns:
The filtered ``DwellSet`` with adjusted cumulative columns and
reduced row count.
"""
logger.info("Filter by dwells by accumulating through")
if not dw.is_dask:
old_len = len(dw.data)
is_long_enough = (
dw.data[params["dwell_time_col"]] >= 0
) # Note, this duration already takes into account the plug in/out times
if params["filter_critical_days"]:
flt_cols = params["filter_cols"]
is_critical = dw.data[flt_cols["refresh"]] | dw.data[flt_cols["crit"]]
is_optional = dw.data[dw.end] <= dw.data[dw.start]
keep_ser = is_long_enough & (is_critical | ~is_optional)
else:
keep_ser = is_long_enough
dw.data["keep_dwells"] = keep_ser
dw = _filter_dwells_core(
dw=dw,
keep_col="keep_dwells",
accum_cols_fw_extra=params["accum_cols_forward_extra"],
accum_cols_rv=params["accum_cols_reverse"],
drop_cols_extra=params["drop_cols"],
)
if not dw.is_dask:
new_len = len(dw.data)
abs_diff = old_len - new_len
pct_diff = round(abs_diff / old_len * 100, 1)
logger.info(f"Rows dropped: {abs_diff}, {pct_diff}%")
else:
logger.info("Rows dropped: not calculated for Dask-backed DwellSets.")
return dw
[docs]
def mark_shift_powers(dw: DwellSet, params: dict) -> DwellSet:
"""Record the maximum charging power available later in the shift at each dwell.
The forward-looking charging-choice algorithm needs to know whether a
vehicle will have access to a high-power charger later in the shift
before it decides whether to charge now. This node pre-computes that
look-ahead value.
The approach:
1. For dwells that are *not* refresh or critical, set their available
power to 0 (they will be auto-skipped by the charging algorithm).
2. Reverse-accumulate the max power with a ``CumAggFunc.MAX`` within
each refresh segment, writing the result to every dwell in the
segment (``write_all=True``).
3. At refresh dwells themselves, overwrite the accumulated value with a
sentinel (``params["final_value"]``), since at a refresh stop the
remaining power budget resets.
4. Shift the accumulated column backward by one position so each dwell
sees the *future* maximum, not its own value.
Args:
dw: Dwell dataset with refresh, critical, and max-power columns
already populated. Modified in-place and returned.
params: Pipeline parameters. Expected keys:
- ``refresh_col`` (str): Boolean column marking refresh-eligible
dwells.
- ``crit_col`` (str): Boolean column marking critical-day dwells.
- ``max_power_col`` (str): Column holding per-dwell maximum
charging power (kW).
- ``final_value`` (float): Sentinel power value assigned to
refresh dwells after accumulation (typically the maximum
possible charger power, kW).
- ``fill_value`` (float): Fill value used for the shift at the
end of a vehicle's sequence (typically 0.0).
- ``max_power_col_shift`` (str): Output column name for the
shifted look-ahead power (kW).
Returns:
The updated ``DwellSet`` with a new look-ahead power column and
intermediate columns removed.
"""
refr_col = params["refresh_col"]
crit_col = params["crit_col"]
max_pow_col = params["max_power_col"]
dont_auto_skip = dw.data[crit_col] | dw.data[refr_col]
dw.data["power_w_skips"] = dw.data[max_pow_col].where(dont_auto_skip, other=0.0)
dw.accum_masked(
keep_mask_col=refr_col,
accum_cols="power_w_skips",
reverse=True,
agg_func=CumAggFunc.MAX,
write_all=True,
inplace=True,
)
acc_col = f"power_w_skips_{refr_col}"
dw.data[acc_col] = dw.data[acc_col].where(
~dw.data[refr_col], other=params["final_value"]
)
kws = {"fill_value": params["fill_value"]}
if dw.is_dask:
kws.update({"meta": ("x", "f8")})
shift_col = params["max_power_col_shift"]
dw.data[shift_col] = dw.data.groupby(dw.veh)[acc_col].shift(-1, **kws)
dw.data = dw.data.drop(columns=["power_w_skips", acc_col])
return dw
[docs]
def simulate_charging_choice(
dw: DwellSet, vehs: pd.DataFrame, modes: pd.DataFrame, params: dict
) -> DwellSet:
"""Run the forward-looking charging-choice simulation for all vehicles.
For each vehicle, the ``ForwardLookingChargingChoiceStrategy`` iterates
through dwells in chronological order and selects — at each stop — the
charging mode and energy amount that maximises a utility function while
respecting battery, power, and delay constraints. The strategy is
inspired by Liu et al. and is implemented with Numba JIT compilation for
speed.
Supports both pandas (single-process) and Dask (distributed) backends.
For Dask, the simulation runs independently on each partition; it is the
caller's responsibility to ensure that each partition contains complete,
sorted vehicle sequences (no vehicle spans multiple partitions).
Optionally pre-compiles the Numba JIT functions using a small mock dataset
before the main run. Pre-compilation is useful for single-process runs
but does not help distributed Dask workers (each worker JIT-compiles on
first use).
Args:
dw: Dwell dataset sorted by vehicle and time. For pandas-backed
DwellSets, sorting is performed automatically; for Dask-backed
DwellSets, the caller must guarantee sort order is preserved
across partitions.
vehs: Vehicle-level table with battery capacity and other per-vehicle
parameters required by the charging-choice strategy.
modes: Charging-modes table produced by ``prepare_modes``, with one
row per mode and columns for power limits and other attributes.
params: Pipeline parameters. Expected keys:
- ``input_cols`` (dict): Column-name mappings forwarded to
``ForwardLookingChargingChoiceStrategy.__init__``. Must
include ``modes_avail`` (bitmask column name) among others.
- ``precompile`` (bool): Whether to trigger Numba pre-compilation
before the main simulation run.
- ``drop_cols`` (list[str]): Columns to drop from ``dw.data``
after the simulation (e.g. intermediate look-ahead columns).
Returns:
The updated ``DwellSet`` with charging-decision columns added (charge
amount, mode chosen, accumulated delay, etc.) and ``drop_cols``
removed.
"""
if dw.is_dask:
logger.info(
"For Dask-backed DwellSets, we assume that sorting has been preserved."
)
else:
dw.sort_by_veh_time()
strat = ForwardLookingChargingChoiceStrategy(**params["input_cols"])
if params["precompile"]: # Note: Pre-compilation does not help distributed workers
logger.info("Pre-compiling charging choice JIT-compiled functions.")
dw_mock = dw.copy_without_data()
base_df = dw.data._meta_nonempty if dw.is_dask else dw.data.iloc[:5].copy()
dw_mock.data = base_df.copy()
# Set all mode bits to 1 for each mock dwell
all_mask = bool_arr_to_bits(np.ones(shape=(len(modes),), dtype=np.bool_))
modes_avail_col = params["input_cols"]["modes_avail"]
if modes_avail_col in dw_mock.data.columns:
dw_mock.data[modes_avail_col] = np.repeat(
all_mask, repeats=len(dw_mock.data)
)
vehs_mock = vehs.iloc[: len(dw_mock.data)].copy()
vehs_mock.index = dw_mock.data.index
_ = strat.run(dwells=dw_mock, vehs=vehs_mock, modes=modes, show_progress=False)
logger.info("Run charging choice simulation.")
if dw.is_dask:
def _run_charging_on_partition(
partition_df: pd.DataFrame, vehs_part: pd.DataFrame, dw_empty: DwellSet
):
# Create a temporary DwellSet for this partition
dw_part = dw_empty.copy_without_data()
dw_part.data = partition_df
result = strat.run(
dwells=dw_part, vehs=vehs_part, modes=modes, show_progress=False
)
return result
# Create meta DataFrame to define output structure using schema generation
meta = strat.get_output_schema(input=dw.data)
vehs_ddf = dd.from_pandas(vehs, npartitions=1)
dw.data = dw.data.map_partitions(
_run_charging_on_partition, vehs_ddf, dw.copy_without_data(), meta=meta
)
else:
dw.data = strat.run(dwells=dw, vehs=vehs, modes=modes)
dw.data = dw.data.drop(columns=params["drop_cols"])
return dw
[docs]
def filter_dwells_post(dw: DwellSet, params: dict) -> DwellSet:
"""Drop optional stops where the vehicle chose not to charge.
Applied *after* the charging-choice simulation. Optional stops are proxy
dwells with zero net duration that were inserted at potential en-route
charging locations (e.g. truck stops along a routed path). If the
simulation determined that no charging occurs at an optional stop, that
row carries no useful information and is removed to keep the output
dataset compact.
A dwell is dropped when both conditions hold:
- It is an optional stop (``dw.end <= dw.start``), AND
- The simulated charge amount is zero.
Non-optional (real) dwells are always kept regardless of charge amount.
When ``filter_unused_optionals`` is ``False``, this function is a no-op.
Args:
dw: Dwell dataset after charging simulation, containing charge-amount
and timing columns. Modified in-place and returned.
params: Pipeline parameters. Expected keys:
- ``filter_unused_optionals`` (bool): Whether to apply the filter.
Set to ``False`` to retain all optional stops (e.g. for
debugging).
- ``filter_cols`` (dict): Required when
``filter_unused_optionals`` is ``True``. Sub-key:
- ``charge`` — column holding the simulated charge amount (kWh)
for each dwell.
- ``accum_cols_forward_extra`` (list[str]): Extra columns to
accumulate in the forward direction through the mask.
- ``accum_cols_reverse`` (list[str]): Columns to accumulate in
the reverse direction through the mask.
- ``drop_cols`` (list[str]): Additional columns to drop after
filtering.
Returns:
The filtered ``DwellSet`` with unused optional stops removed and
cumulative columns corrected.
"""
logger.info("Filter by dwells by accumulating through")
if not dw.is_dask:
old_len = len(dw.data)
if params["filter_unused_optionals"]:
is_optional = dw.data[dw.end] <= dw.data[dw.start]
has_some_charge = dw.data[params["filter_cols"]["charge"]] > 0.0
dw.data["keep_dwells"] = ~is_optional | (is_optional & has_some_charge)
dw = _filter_dwells_core(
dw=dw,
keep_col="keep_dwells",
accum_cols_fw_extra=params["accum_cols_forward_extra"],
accum_cols_rv=params["accum_cols_reverse"],
drop_cols_extra=params["drop_cols"],
)
else:
pass
if not dw.is_dask:
new_len = len(dw.data)
abs_diff = old_len - new_len
pct_diff = round(abs_diff / old_len * 100, 1)
logger.info(f"Rows dropped: {abs_diff}, {pct_diff}%")
else:
logger.info("Rows dropped: not calculated for Dask-backed DwellSets.")
return dw