Creating Steps with Parameters

This page explains how to add a new Step to the Pymetropolis codebase. It assumes you are already familiar with the Pipeline (Steps, MetroFiles, Parameters) as a user of Pymetropolis, and, for the Creating a PopulationStep section, with the Multiple-demand concept page.

Basics

A Step is declared as a Python class, in a steps.py (or similarly-named) module of the relevant package (e.g. src/pymetropolis/metro_demand/population/eqasim.py), subclassing Step.

from pymetropolis.metro_pipeline import Step

from .files import TripsDestinationsFile, TripsDistancesFile, TripsOriginsFile


class TripDistancesStep(Step):
    """Computes the Euclidean distances between origin and destination for each trip."""

    input_files = {"origins": TripsOriginsFile, "destinations": TripsDestinationsFile}
    output_files = {"distances": TripsDistancesFile}

    def run(self):
        import polars as pl

        origins = self.input["origins"].read()
        destinations = self.input["destinations"].read()
        distances = pl.DataFrame(
            {"trip_id": origins["trip_id"], "od_distance": origins.distance(destinations)}
        )
        self.output["distances"].write(distances)

• input_files/output_files are dict[str, type[MetroFile]]: keys are arbitrary names used only to look files up inside run() (self.input["origins"]/self.output["distances"]), values are the MetroFile subclasses the Step reads/writes. The Pipeline uses these declarations, not the body of run(), to build the dependency graph — always keep them in sync with what run() actually reads and writes, or caching will be wrong.
• The class docstring becomes the Step’s description in the generated Steps reference — write one for any Step whose behavior isn’t obvious from its name.
• run() is where the Step’s actual logic lives. Import heavy third-party packages (polars, geopandas, duckdb, …) inside run() rather than at module level, to keep the CLI’s startup time fast — every Step module is imported once at startup to build the pipeline’s step list, even for Steps that never actually run.

Once your class is defined, add it to the STEPS list exported by your package’s __init__.py (see how POPULATION_STEPS in src/pymetropolis/metro_demand/population/__init__.py is aggregated), so MetroPipeline discovers it. A Step that is never added to a STEPS list is simply invisible to Pymetropolis: it neither runs nor appears in the documentation.

Optional and conditional input files

A bare MetroFile class in input_files means the input is required and always needed. Wrap it in InputFile for anything else:

from pymetropolis.metro_pipeline.steps import InputFile

input_files = {
    "simulation_area": InputFile(SimulationAreaFile, optional=True),
    "uniform_draws": InputFile(
        UniformDrawsFile,
        when=lambda step: step.has_mode_choice(),
        when_doc="if there are at least two modes",
    ),
}

• optional=True: the Step can run even if the file was never produced (self.input["name"] still resolves to a MetroFile instance, but reading it may fail if it truly doesn’t exist on disk — guard with .exists() first, or handle the missing case, as needed).
• when=lambda step: ...: the input is only demanded when the callable returns True for this Step instance (evaluated after the Step’s parameters are resolved, so it can inspect self/step’s parameter values, as in the mode-choice example above). Pass when_doc alongside it — a short, human-readable explanation shown in the generated reference next to when, since the lambda itself isn’t rendered there.

is_defined and priority

Override is_defined(self) -> bool (default: always True) to have the Step skipped entirely rather than run when its required configuration is missing — see EqasimImportStep.is_defined in src/pymetropolis/metro_demand/population/eqasim.py, which returns False when synthetic_population.eqasim_output isn’t set, so a config that uses a different population source doesn’t fail on this Step’s other, unrelated required parameters.

Override the priority: ClassVar[int] class attribute (default 1) to mark a non-primary Step (priority = 0): it only runs if another primary Step actually needs its output, rather than unconditionally whenever its inputs are available. Use this for Steps with potentially long running time whose outputs are not useful per se (e.g., AllRoadDistancesStep which computes the shortest path distance for all node pairs in the road network).

Parameters

A Parameter binds a config TOML key to a typed, validated attribute on the Step, declared as a class attribute:

from pymetropolis.metro_pipeline.parameters import FractionParameter, PathParameter


class EqasimImportStep(Step):
    # Declared as class attributes, as `Parameter` instances.
    eqasim_output = PathParameter(
        "synthetic_population.eqasim_output",
        check_dir_exists=True,
        description="Path to the output directory of the Eqasim synthetic population pipeline.",
    )
    fraction = FractionParameter(
        "synthetic_population.fraction",
        default=1.0,
        description="Fraction of the synthetic population to be selected for simulations.",
        note=(
            "If the synthetic population already represents a part of the total population, you "
            "probably want to keep this parameter to 1 and instead set `simulation_ratio` to the "
            "actual share of the population being simulated."
        ),
    )

    def is_defined(self):
        # Only run the Step if the `eqasim_output` parameter is set in the config (i.e., it is not
        # equal to its default value (`None`).
        return self.eqasim_output is not None

    def run(self):
        # The parameters directly resolve as their validated value (a `pathlib.Path` and a `float`
        # here), not as the `Parameter` objects declared above.
        assert isinstance(self.eqasim_output, Path)
        assert isinstance(self.fraction, float)
        assert self.fraction <= 1.0
        ...

Even though eqasim_output/fraction are declared as Parameter instances at the class level, self.eqasim_output/self.fraction never evaluate to those Parameter objects on an actual Step instance: once the Step is constructed, they resolve to the validated config value itself (or None/the default when the TOML key is absent) — a pathlib.Path, a float, and so on, depending on the Parameter subclass. This holds anywhere on the instance, not just in run(): in is_defined() above, self.eqasim_output is already the resolved Path | None, which is exactly what makes checking it there a meaningful “is this Step configured at all” test.

Every Parameter subclass shares these constructor arguments:

• key (positional, first argument): the dotted TOML key, e.g. "synthetic_population.fraction".
• default: value used when the key is absent from the config. Left as None (meaning “no default”, i.e. the attribute resolves to None) if omitted.
• description, example, note: free-text, all optional, used to build the generated Parameters reference. note is for caveats or cross-references to other parameters (see fraction above); example is for a config snippet, mainly useful for parameters whose valid values aren’t obvious from description alone.

Pick the typed subclass that matches the value, rather than the base Parameter:

random.py and metro_spatial/crs.py add further examples of CustomParameter (GeoStep.crs, validating/normalizing a CRS with pyproj) and distribution-valued parameters (FloatDistributionParameter, etc., accepting either a constant or a {mean, std, distribution} table).

Shared parameters

Pass shared=True for a Parameter that is genuinely global rather than per-population. In a PopulationStep (see below), parameter values are read from the population-specific config, except for shared parameters where the value in the main configuration is used as a fallback when the parameter is not defined in the population-specific config. This is useful for the random_seed parameter, for example, so that it needs to be defined only once in the main config.

Step inheritance

A Step can subclass another Step (in addition to Step itself) to reuse its Parameters and any helper methods, without repeating them. RandomStep (pymetropolis/random.py) adds self.random_seed plus a get_rng(step_name) helper; GeoStep (metro_spatial/crs.py) adds self.crs. A Step can mix in several of these at once:

from pymetropolis.metro_pipeline import PopulationStep
from pymetropolis.metro_spatial import GeoStep
from pymetropolis.random import RandomStep


class EqasimImportStep(GeoStep, RandomStep, PopulationStep):
    ...

    def run(self):
        ...
        homes = read_homes(..., random_seed=self.random_seed)
        homes = homes.to_crs(self.crs)
        ...

When a Step needs its own random number generator (rather than just passing self.random_seed through to a helper, as above), call self.get_rng(str(self)) — always with the calling Step’s own str(self), never a literal string or another Step’s — so that Steps/populations sharing random_seed still draw independent random sequences from one another.

Creating a PopulationStep

If your Step should run once per configured population (persons, trucks, …) rather than once per pipeline run — the common case for anything on the demand side — subclass PopulationStep instead of Step. Everything above (Parameters, InputFile, is_defined, mixins, …) still applies identically; the differences are:

• Every non-optional output_files entry must be a PopulationFile — see Creating MetroFiles — since a plain MetroFile output would resolve to the identical path for every population and silently discard all but one population’s result. This is enforced at class-definition time: subclassing PopulationStep with a non-PopulationFile output raises immediately.
• self.population_name holds which population this particular instance is running for (the main config’s sentinel "population", or an extra population’s own name).
• str(self) — used for progress-log output and to key the on-disk cache file (main_directory/update_files/<str(self)>.json, see the note above) — is the population name and the class name joined with a double underscore for every population but the main one, e.g. trucks__TripDistancesStep (plain TripDistancesStep for the main population). This is also why get_rng(str(self)) below draws an independent random sequence per population, even when they share the same random_seed.
• Non-shared Parameters resolve against that population’s own config (falling back to the main config only when shared=True — see Shared parameters)
• PopulationFile inputs/outputs resolve against that population’s namespaced path automatically — self.input/ self.output work exactly as for a plain Step, you don’t need to pass the population name around yourself.

from pymetropolis.metro_pipeline import PopulationStep

from .files import TripsDestinationsFile, TripsDistancesFile, TripsOriginsFile  # all PopulationFile


class TripDistancesStep(PopulationStep):
    """Computes the Euclidean distances between origin and destination for each trip."""

    input_files = {"origins": TripsOriginsFile, "destinations": TripsDestinationsFile}
    output_files = {"distances": TripsDistancesFile}

    def run(self):
        # Identical body to the plain-Step example above: `self.input`/`self.output` are already
        # resolved to this instance's own population.
        # The step will automatically run once per population.
        ...

Reading every population at once

A plain (non-PopulationStep) Step that needs to consume every population’s copy of a PopulationFile at once — typically to merge them into a single file for actual METROPOLIS2 input — declares that input with InputFile(..., all_populations=True) and reads it through self.input_populations["name"] (a dict[population_name, MetroFile]) instead of self.input:

from pymetropolis.metro_pipeline import Step
from pymetropolis.metro_pipeline.steps import InputFile
from pymetropolis.metro_simulation.common import merge_populations

from .files import MetroAgentsFile, MetroAgentsPopulationFile


class WriteMetroAgentsStep(Step):
    """Merges the agents in each population and writes the agents input file for Metropolis-Core."""

    input_files = {"population_agents": InputFile(MetroAgentsPopulationFile, all_populations=True)}
    output_files = {"metro_agents": MetroAgentsFile}

    def run(self):
        dfs = [
            f.read().with_columns(
                agent_id=pl.concat_str(pl.lit(f"{population}-"), pl.col("agent_id"))
            )
            for population, f in self.input_populations["population_agents"].items()
        ]
        agents = pl.concat(dfs, how="vertical")
        self.output["metro_agents"].write(agents)

By default (optional=False) such a Step only becomes feasible if all configured populations actually produce the file; pass optional=True on the InputFile if the Step should run regardless of how many populations produced it (there is no built-in “at least one, but not necessarily all” option). merge_populations (metro_simulation/common.py) is a ready-made helper for the common case of concatenating every population’s rows into one DataFrame, prefixing the given id columns with the population name so ids stay globally unique after the merge — reuse it rather than writing the same concatenation logic again when your merge is a straightforward row-wise union.

Note

all_populations=True requires a PopulationFile and is only meaningful on a non-PopulationStep. A plain (non-all_populations) PopulationFile input on a non-PopulationStep, or any PopulationFile input/output at all on a Step that is neither a PopulationStep nor declared all_populations=True, is rejected at class-definition time — see the two __init_subclass__ checks in steps.py — since either would silently only ever see the main population’s copy.