Add docstrings to facilitate newcomers getting started

[MatKbauer]

What maintenance does this project need?

Some classes and functions are hard to understand for newcomers due to sparse documentation and missing docstrings, in particular. Improvements can be made to the following classes.

model.py: The predict function of the Model class could be confusing, as it rather decodes from latent to physical space (thanks to @kctezcan for explaining) instead of generating a prediction. On the other hand, the decoder fills gaps, which are predictions then. Any thoughts on how to make this more approachable appreciated. A corresponding docstring should resolve this unclarity or predict can be renamed in decode.

engines.py: Write docstrings to describe more exhaustively what each engine is doing. I.e., what's the difference between local and global assimilation engine?

attention.py: Add docstrings and describe the purpose of the different attention heads (local, varlen, flexed, slicedq)

Organisation

ECMWF

[ankitpatnala]

#268

@MatKbauer Can you also give a look to this pull request by @TillHae. Just to avoid repeating things.

[mlangguth89]

I strongely appreciate the suggestion to rename predict into decode. Actually, to predict and to forecast are synonyms.

[clessig]

I strongely appreciate the suggestion to rename predict into decode. Actually, to predict and to forecast are synonyms.

Prediction head is standard terminology in machine learning. An LLM prediction head does "forecasting" of the next token but I think "forecasting head" is not a good alternative or synonymous to "prediction head".

[MatKbauer]

Agree that prediction_head is standard terminology; the name of the module seems fine to me. It remains confusing to me, though, to have a predict and a forecast function next to each other. More suggestions to resolve ambiguities:

1. Rename predict into diagnose to distinguish from temporal prediction (forecasting) and to emphasize the within-one-time-step process.
2. Separate (A) the prediction of masked tokens (gap filling) from (B) the decoding into physical space. Step (A) can be named predict_masked_tokens or just infer_masked and step (B) could become decode_latent_to_observations, decode, or reconstruct.