Custom Steps

Pymetropolis ships with a large library of built-in Steps, but some data sources are inherently local: a country- or region-specific household travel survey, a proprietary origin-destination matrix, a dataset with its own idiosyncratic format. It rarely makes sense to add a Step for these to the Pymetropolis codebase itself, since it would only ever be useful to a handful of users — yet you still want it to participate in the pipeline like any other Step: cached, ordered, and wired up automatically based on the files it reads and writes.

The custom_steps parameter lets you define your own Step subclasses in a plain Python file and have Pymetropolis load them alongside its built-in Steps.

Tip

For a complete, runnable config, see examples/extra/bottleneck-custom-steps/ in the Pymetropolis repository: the Bottleneck case study extended with three custom Steps, each demonstrating a different extension pattern.

Declaring a custom steps file

# config.toml
custom_steps = ["my_steps.py"]

• Allowed values: list of strings representing valid paths.
• Paths are resolved relative to the main configuration file.
• Several files can be listed; each one is loaded independently.

my_steps.py is an ordinary Python module. Anything you can do when creating a built-in Step — Step/PopulationStep, Parameters, MetroFiles, RandomStep, GeoStep, etc. — works exactly the same way here. A custom Step can also target a built-in MetroFile as its output, not just a new one of its own — this is how it plugs into the rest of the pipeline, exactly like a built-in Step would. For example, CircularNetworkStep and OpenStreetMapRoadImportStep both produce RoadEdgesRawFile, the road-network entry point consumed by all of the built-in road-network cleanup Steps downstream; a custom Step producing that same file from our own local dataset plugs into that same downstream pipeline for free:

# my_steps.py
from pymetropolis.metro_common.io import read_geodataframe
from pymetropolis.metro_network.road_network import RoadEdgesRawFile
from pymetropolis.metro_pipeline import Step
from pymetropolis.metro_pipeline.parameters import PathParameter


class LocalRoadNetworkImportStep(Step):
    """Converts our local road-network dataset into Pymetropolis's raw road-edges format."""

    network_path = PathParameter(
        "local_road_network.path",
        check_file_exists=True,
        description="Path to the local road-network dataset.",
    )
    output_files = {"raw_edges": RoadEdgesRawFile}

    def is_defined(self) -> bool:
        return self.network_path is not None

    def run(self):
        edges = read_geodataframe(self.network_path)
        # ... derive the `edge_id`, `source`, `target`, `length`, and `speed_limit` columns
        # `RoadEdgesRawFile` expects from our local dataset's own columns ...
        self.output["raw_edges"].write(edges)

is_defined keeps the Step out of the pipeline entirely (rather than failing) when local_road_network.path isn’t set, e.g. in a config that gets its road network from OpenStreetMap or the built-in circular toy network instead — see is_defined and priority.

Every Step subclass defined directly in the file (not merely imported into it) is picked up automatically — there is nothing else to register. See Creating Steps with Parameters for the full Step/Parameter/MetroFile API: it is identical whether the Step lives in the Pymetropolis codebase or in your own custom_steps file.

Overriding a built-in Step

If a custom Step’s class name matches an existing Step’s name — built-in, or from another custom_steps file — it replaces it rather than causing an error. This is intended: it lets you swap out a single built-in Step for your own local-specific implementation.

Since it is still an ordinary Python class, the most common way to do this is to subclass the real built-in Step. This gives you everything the built-in Step already declares (input_files, output_files, is_defined, existing Parameters, mixed-in helpers like RandomStep.get_rng) for free, so you only need to write the part that’s actually different.

For example, say we want fuel consumption to also depend on free-flow travel time, not just free-flow distance — to account for the extra fuel burned idling and accelerating in slower traffic, on top of the existing per-km rate. CarFuelStep already reads TripsCarFreeFlowTravelTimesFile, which has both a free_flow_distance and a free_flow_travel_time column — run() just doesn’t use the latter. So the override doesn’t even need to touch input_files: only one new Parameter and run():

# my_steps.py
from pymetropolis.metro_common.utils import pl_duration_to_seconds
from pymetropolis.metro_environment.fuel import CarFuelStep
from pymetropolis.metro_pipeline.parameters import FloatParameter


class CarFuelStep(CarFuelStep):
    """Same as the built-in Step, but adds a per-hour fuel-consumption term on top of the
    per-km one, to account for the extra fuel burned idling/accelerating in slower traffic.
    """

    fuel_factor_time = FloatParameter(
        "fuel.consumption_factor_time",
        default=0.0,
        description="Extra fuel consumption per hour of free-flow travel time, in liters per hour.",
    )

    def run(self):
        import polars as pl

        df = self.input["ff_distances"].read()
        df = df.select(
            "trip_id",
            fuel_consumption=self.fuel_factor * pl.col("free_flow_distance") / 1000.0
            + self.fuel_factor_time * pl_duration_to_seconds("free_flow_travel_time") / 3600.0,
        )
        # If `fuel_price` is not defined (None), `fuel_cost` will be all null values.
        df = df.with_columns(fuel_cost=pl.col("fuel_consumption") * self.fuel_price)
        self.output["fuel_consumption"].write(df)

input_files, output_files, is_defined, and the existing fuel_factor/fuel_price Parameters are all inherited unchanged — only fuel_factor_time and run() are new. The function pl_duration_to_seconds is an helper function from the built-in codebase to turn a Duration column like free_flow_travel_time into a number of seconds. Every downstream Step still only depends on CarFuelFile (not on CarFuelStep itself), so the rest of the pipeline runs unmodified against whatever this override produces.

Note

The class name is also what keys the on-disk cache (main_directory/update_files/<ClassName>.json), so overriding a Step this way reuses its cache file. Declare every config value your run() actually reads as a Parameter on your class — here, fuel_factor_time — exactly as for any other Step, or Pymetropolis won’t know to invalidate the cache when it changes.