Tracking of Input Channel Indices <-> Variable Name and Level

[sadamov]

Description

Tracking of which feature-channel corresponds to which variable and vertical level of the input data is benefitial for plotting, verification and more...

1. precompute_variable_indices() method to precompute the indices for each variable in the input tensor.
2. selected_vars_units attribute to store the short names and units of the selected variables.

Implementation

Suggestion to add the following code to the ARModel class:

self.variable_indices = self.precompute_variable_indices()
self.selected_vars_units = list(
    zip(constants.PARAM_NAMES_SHORT, constants.PARAM_UNITS)
)

def precompute_variable_indices(self):
"""
Precompute indices for each variable in the input tensor
"""
variable_indices = {}
all_vars = []
index = 0
# Create a list of tuples for all variables, using level 0 for 2D
# variables
for var_name in constants.PARAM_NAMES_SHORT:
    if constants.IS_3D[var_name]:
        for level in constants.VERTICAL_LEVELS:
            all_vars.append((var_name, level))
    else:
        all_vars.append((var_name, 0))  # Use level 0 for 2D variables

# Sort the variables based on the tuples
sorted_vars = sorted(all_vars)

for var in sorted_vars:
    var_name, level = var
    if var_name not in variable_indices:
        variable_indices[var_name] = []
    variable_indices[var_name].append(index)
    index += 1

return variable_indices

Benefits

• The precompute_variable_indices() method will allow efficient lookup of indices for each variable in the input tensor, avoiding the need to recompute them during runtime. This is very flexible as the user can define it in constants.py
• The selected_vars_units attribute will provide easy access to the short names and units of the selected variables, which can be useful for plotting, logging, or other purposes.

[joeloskarsson]

I agree that this is very useful to have. I would rather extend the Config class with it than put it as part of the ARModel. As mentioned above, this is very useful information that is needed in many places in the code. Therefore I would not want to associate it with the model class specifically. These maps are also direct consequences of the data config, making it logical to access from the same object.

Maybe it would be worth to wait with this until after #24 is done? But I would not be against introducing this now and just adapting it when we change that.