Improve output_callback data recognition

[HCookie]

As output_callback is called for input_fields, prognostic and diagnostics, it may be unclear for a provided function what data is being passed through.

Current Solution

Examples using anemoi-inference look for steps in the kwargs

if 'steps' in kwargs:
   do_things()

Proposed Solution

• Data class including information about source and type of data

class ModelDataType(Enum):
    """Model Data Type Enum"""

    PROGNOSTIC = 0
    DIAGNOSTIC = 1
    FORCING = 2

class ModelDataSource(Enum):
    """Model Data Source Enum"""

    INPUT = 0
    PREDICTION = 1

@dataclass
class ModelDataInfo():
    """Model Data Info Dataclass"""

    type: ModelDataType
    source: ModelDataSource

    def __str__(self):
        return f"{self.type} {self.source}"

Usage:

model_info = ModelDataInfo(ModelDataType.PROGNOSTIC, ModelDataSource.PREDICTION)
if ModelDataSource.INPUT = model_info.source:
    do_other_things()

Output_callback

def output_callback(data, model_info: ModelDataInfo, *args, **kwargs):
    ...

Looking for thoughts

• Will this limit the output callback too much?
• Will this scale well?