nvalchemi.models.pipeline.PipelineStep

class nvalchemi.models.pipeline.PipelineStep(model, wire=<factory>)

Wraps a model with an output rename mapping.

Only needed when a model’s output key doesn’t match the downstream input key. For models that don’t need renaming, pass the bare model directly — the pipeline normalizes it internally.

Parameters:

• model (BaseModelMixin) – The model to wrap.

• wire (dict[str, str]) – Output-to-attribute rename mapping. Each entry {output_key: data_attribute} causes the pipeline to write the model’s output_key value onto data.data_attribute before downstream models execute. Downstream models that declare data_attribute in their required_inputs will then receive it automatically.

Examples

AIMNet2 produces "charges" (per-atom partial charges), but the Ewald model expects "node_charges" as a required input:

PipelineStep(aimnet2, wire={"charges": "node_charges"})

After AIMNet2 runs, the pipeline writes its "charges" output onto data.node_charges. When Ewald runs next, its adapt_input() finds data.node_charges and uses it.

If a model’s output keys already match downstream input keys, no wire mapping is needed — pass the bare model:

PipelineGroup(steps=[model_a, model_b])  # auto-wired