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

A wire mapping is only needed when a producer’s output key differs from the consumer’s required-input key. For example, if a charge model emits "charges" but a downstream model requires them under "partial_charges":

PipelineStep(charge_model, wire={"charges": "partial_charges"})

After the charge model runs, the pipeline writes its "charges" output onto data.partial_charges. When the downstream model runs next, its adapt_input() finds data.partial_charges and uses it.

When a model’s output keys already match downstream input keys, no wire mapping is needed — pass the bare model. AIMNet2 outputs "charges" and EwaldModelWrapper requires "charges", so the two auto-wire on that key:

PipelineGroup(steps=[aimnet2, ewald])  # auto-wired on "charges"