Automatically derive some model params (and remove from config)

[charlie-becker]

Our current config has a few parameters that relate to the data input and output channels that can be easily derived from the new data structure itself.

model:
    type: "wxformer"
    frames: 1
    image_height: 640
    image_width: 1280
    levels: 16
    channels: 4
    surface_channels: 1
    input_only_channels: 0
    output_only_channels: 0   # no diagnostics for speed

Specifically, levels, channels, surface_channels, input_only_channels and output_only_channels can be directly derived from the data portion of the config [prognostic, diagnostic, dynamic_forcing, and static.

I propose removing those entirely from the config and calculating them from the data config within the model init step which should work for both new models and loading in pretrained models.

Furthermore, it would be nice to do this with image_height, image_width and padding too, but that would need to be derived from the data shape of the tensors (and not just the config) and then would also have to make some mathematical assumptions about the padding sizes, but this would require more thought.

[djgagne]

I do think there might be a way to do this as a wxformer-simple model wrapper either on top of regular wxformer or importing wxformer components. If we can initialize all model components on the first forward call with a given dataset rather than in ___init__, then we should be able to figure out data shapes. Where it may be trickier is loading in pretrained model weights.

[djgagne]

Some considerations for a more robust model class that need to be addressed:

• Current model loading process for PyTorch involves creating a new instance of a model class, initializing the weight matrices, and then loading weights from the pkl file into the the initialized weight matrices.
• Current saving of model.yaml only saves what parameters are in the config file, so if default keyword arguments are changed for a model class, then the reloaded model might have different behavior from before. We need to have a way to save both specified and default hyperparameters to save_loc.
• We want to preserve backwards compatibility for existing WXFormer models, and we want to minimize reuse of code. New model class can either import or copy code from other modules. Copying is safest to minimize reproducibility issues.