Extend support for manual specification of MetricCollection compute_groups

[alexrgilbert]

🚀 Feature

Support incomplete specification of MetricCollection compute_groups + inherit compute groups from child metric collections

Motivation

Currently, new instances of the MetricCollection class allow users to manually specify the compute_groups used to simplify update steps. However, if a user doesn't explicitly specify any metric in the compute groups lists, it simply won't be updated. It seems like the expected behavior would be for those additional metrics to simply be updated normally (or possibly automatically inherited by some of the manually specified groups).

Moreover, if a uses the supported behavior of a single level "recursion" of MetricCollection's (i.e. creating a MetricCollection of MetricCollection's), the compute_groups arguments of the child collections are simply ignored.

Pitch

1. It might be desirable to have any metrics which aren't included in the list of compute_groups to be automatically detected and put into their own single-item compute groups. Possibly we could also attempt to merge them into existing groups, but this might be overly complex and unnecessary.

2. The specified compute groups from child collections should be inherited by the parent (with the prefix's/postfix's updated appropriately). Admittedly, this introduces some complexities which may be more simply handled by just automatically creating the groups (as is currently done). If so, it should at least be documented that the child compute_groups arguments will be ignored.

Alternatives

As noted above, if the current behavior is intended, it should at least be documented what will happen in these "corner" cases.

Additional context

[SkafteNicki]

Hi @alexrgilbert, thanks for reporting this issue.
I created PR #2979 which will fix the first issue you mention, namely that if a metric is not in the list of compute groups it will be automatically added still. For you second issue, I really tried to find a solution for that, but was unable for now to get it working because there are so many parts moving inside the metric collection. For now, i post my part-solution in this issue such that I can hopefully fix it in the future and in addition added this limitation to the description of the metric collection.

# add inside the _init_compute_groups() method
    # if nested metric collections, we need to respect the compute groups of child collections
    if self._groups:
        
        flattened_dict = {}
        counter = 0
        for outer_dict in self._groups.values():
            if isinstance(outer_dict, dict):
                for values in outer_dict.values():
                    flattened_dict[counter] = values
                    counter += 1  # Increment the unique key
            else:
                flattened_dict[counter] = outer_dict
                counter += 1
        self._groups = flattened_dict

        for k, v in self.items(keep_base=False, copy_state=False):
            if hasattr(v, "_from_collection"):
                if hasattr(v, "_collection_name"):
                    kmatching = k.lstrip(v._collection_name)
                    for group_key, values in self._groups.items():
                        self._groups[group_key] = [k if val == kmatching else val for val in values]
                else:
                    for group_key, values in self._groups.items():
                        self._groups[group_key] = [k if val == k else val for val in values]

    # add metrics not already added in group by child collections
    existing_metrics = {metric for values in self._groups.values() for metric in values}
    new_metrics = [k for k in self.keys(keep_base=True) if str(k) not in existing_metrics]
    for i, metric in enumerate(new_metrics, start=len(self._groups)):
        self._groups[i] = [str(metric)]  
``