5. Using STRUCTS¶

This tutorial is based on example code which can be found in the TRAC GitHub Repository under examples/models/python.

Not all data is naturally tabular. Structured data lets a model define an input or output as a Python dataclass (or Pydantic model), including nested objects, enums and dictionaries, rather than a flat table of fields.

Defining a struct schema¶

A struct schema is built from an ordinary Python dataclass. Structs can be nested inside one another and can contain enums, dictionaries and other common Python types.

src/tutorial/structured_objects.py¶

import enum

import typing as _tp
import dataclasses as _dc
import datetime as _dt

import tracdap.rt.api as trac


class EvolutionModel(enum.Enum):
    PERTURB = 1
    SCATTER = 2
    STOCHASTIC = 3

@_dc.dataclass
class ScenarioConfig:

    scenario_name: str
    default_weight: float
    evolution_model: EvolutionModel
    apply_smoothing: bool

@_dc.dataclass
class RunConfig:

    include_front_book: bool
    base_date: _dt.date

    base_scenario: ScenarioConfig
    stress_scenarios: dict[str, ScenarioConfig]

To turn a dataclass into a schema, use define_struct(). The runtime validates the supplied type and builds a matching SchemaDefinition, which can be used with define_input() in the usual way. For outputs, the shorthand define_output_struct() combines both steps.

    def define_inputs(self) -> _tp.Dict[str, trac.ModelInputSchema]:

        run_config_struct = trac.define_struct(RunConfig)
        run_config = trac.define_input(run_config_struct, label="Run configuration")

        return {"run_config": run_config}

    def define_outputs(self) -> _tp.Dict[str, trac.ModelOutputSchema]:

        modified_config = trac.define_output_struct(RunConfig, label="Modified config for next model stage")
        return {"modified_config": modified_config}

Reading and writing structured data¶

Structured inputs and outputs are read and written with get_struct() and put_struct(), passing the dataclass type to use for the result. The runtime returns (and validates) an instance of that type, so the rest of the model code can work with ordinary Python objects and attribute access, rather than looking up fields by name.

    def run_model(self, ctx: trac.TracContext):

        run_config = ctx.get_struct("run_config", RunConfig)

        new_scenario = ScenarioConfig(
            scenario_name="hpi_shock",
            default_weight=1.0,
            evolution_model=EvolutionModel.STOCHASTIC,
            apply_smoothing=True)

        run_config.stress_scenarios["hpi_shock"] = new_scenario

        ctx.put_struct("modified_config", run_config)

Note

The type passed to get_struct() does not have to be the exact type used in define_inputs(), so long as the schema is compatible the runtime will perform the conversion. In practice, models should normally use the same type in both places, as this example does with RunConfig.

This model reads a run configuration struct, adds a new stress scenario to it, and saves the modified configuration as an output - which could then be picked up as the input to another model further down a flow.