Variable Metadata Extractor not working as expected

[JPXKQX]

What happened?

The ExtractVariableGroupAndLevel class allows the user to specify a shortcut (e.g. z) for including all pressure levels (z_50, z_100, z_500, ...). This shortcut can match the name of a surface variable (z). This results in inconsistent behaviour that will trigger an error in anemoi-training when pressure level scaling is used for that variable.

What are the steps to reproduce the bug?

from anemoi.training.utils.variables_metadata import ExtractVariableGroupAndLevel
from anemoi.datasets import open_dataset

ds = open_dataset("aifs-ea-an-oper-0001-mars-o96-1979-2023-6h-v8")

config = {
    "default": "sfc",
    "pl": {"param": ["q", "t", "u", "v", "w", "z"]},
    "sfc": {"param": ["tp", "cp", "10u", "10v", "2d", "sp", "z"]}
}
extractor = ExtractVariableGroupAndLevel(config, ds.metadata()["variables_metadata"])
print(f"Variable: z, Group: {extractor.get_group('z')}")
print(f"Variable: t, Group: {extractor.get_group('t')}")

Version

main

Relevant log output

Variable: z, Group: pl
Variable: t, Group: pl

Organisation

ECWMF

[caglarkucuk]

Just giving it a shot:
i) Making the variables 'complex' by adding an is_surface_level flag in the config:

config = {
    "default": "sfc",
    "pl": {"param": ["q", "t", "u", "v", "w", "z"], "is_surface_level": False},
    "sfc": {"param": ["tp", "cp", "10u", "10v", "2d", "sp", "z"], "is_surface_level": True}
}

and
ii) Changing the 'without metadata' logic of the 'complex group' in the get_group method in ExtractVariableGroupAndLevel class (Lines 116-132) to:

                    # If metadata is not available, fall back to param-only matching
                    if "param" in group_spec:
                        if self.get_param(variable_name) in (
                            group_spec["param"] if isinstance(group_spec["param"], list) else [group_spec["param"]]
                        ):
                            LOG.debug(
                                "Variable %r is in group %r through param-only matching (no metadata available).",
                                variable_name,
                                group_name,
                            )
                            return group_name

yields

Variable: z, Group: sfc
Variable: t, Group: pl

But I'm not sure if that'd affect other functionalities, and it wouldn't be the most elegant solution to escalate group_spec of variables into a 'complex' group. Nevertheless, let me know if this approach makes sense, and I can create a PR with this proposal.

[JPXKQX]

Hi Caglar,
My preferred option would be to support extra keyword arguments that are passed through the selector. We should therefore be able to pass "levtype": 'pl', similar to your 'is_surface_level'. False. This way, a user could use 'pl500': {'level': 500} to select all variables at level 500. That said, I don't know much about this class or all of its use cases, so I'm not sure what the best option is.

Additionally, I would expect extractor.get_group("t") to fail, as "t" is only a shortcut used in the config to select all pressure levels of a variable, not a variable itself. Therefore, I would expect the code to fail if we specify a variable name that is not in the dataset.