Repository: Project-MONAI/MONAI
File: monai/transforms/compose.py (Python)

Code before:
    At each invocation the set of transforms to apply is uniformly sampled.

    Args:
        transforms: list of callables.
        map_items: whether to apply transform to each item in the input `data` if `data` is a list or tuple.
            Defaults to `True`.
        unpack_items: whether to unpack input `data` with `*` as parameters for the callable function of transform.
            Defaults to `False`.
        log_stats: whether to log the detailed information of data and applied transform when error happened,
            for NumPy array and PyTorch Tensor, log the data shape and value range,
            for other metadata, log the values directly. Default to `False`.
        min_num_transforms: minimum number of transforms to sample. Defaults to `0`.
        max_num_transforms: maximum number of transforms to sample. Defaults to `3`.
        fixed: whether to sample exactly `max_num_transforms` transforms, or up to it. Defaults to `False`.
        replace: whether to sample with replacement. Defaults to `False`.
        weights: weights to use in for sampling transforms. Will be normalized to 1. Default: None (uniform).
    """

    def __init__(
        self,
        transforms: Sequence[Callable] | Callable | None = None,
        map_items: bool = True,
        unpack_items: bool = False,
        log_stats: bool = False,
        min_num_transforms: int = 0,
        max_num_transforms: int = 3,
        fixed: bool = True,
        replace: bool = False,
        weights: list[int] | None = None,
    ) -> None:
        super().__init__(transforms, map_items, unpack_items, log_stats)
        if transforms is None:
            transforms = []
        self.n_transforms = len(transforms)
        self.min_num_transforms = min_num_transforms
        self.max_num_transforms = min(self.n_transforms, max_num_transforms)
        self.fixed = fixed
        self.replace = replace
        self.weights = self._normalize_probabilities(weights)

    # Modified from OneOf
    def _normalize_probabilities(self, weights):
        if weights is None or self.n_transforms == 0:
            return None

        weights = np.array(weights)

        n_weights = len(weights)
        if n_weights != self.n_transforms:
            raise ValueError(
                f"The number of weights specified must be equal to the number of transforms provided. Expected: {self.n_transforms}, got: {n_weights}."

Review comment (suggestion):
I'd suggest setting ```max_num_transforms``` to ```None```. When it is ```None```, it can internally set ```max_num_transforms``` to ```len(transforms)```

Code after:
            return weights
        weights = np.array(weights)
        if np.any(weights < 0):
            raise ValueError(f"Probabilities must be greater than or equal to zero, got {weights}.")
        if np.all(weights == 0):
            raise ValueError(f"At least one probability must be greater than zero, got {weights}.")
        weights = weights / weights.sum()
        return list(weights)

    def flatten(self):
        transforms = []
        weights = []
        for t, w in zip(self.transforms, self.weights):
            # if nested, probability is the current weight multiplied by the nested weights,
            # and so on recursively
            if isinstance(t, OneOf):
                tr = t.flatten()
                for t_, w_ in zip(tr.transforms, tr.weights):
                    transforms.append(t_)
                    weights.append(w_ * w)
            else:
                transforms.append(t)
                weights.append(w)
        return OneOf(transforms, weights, self.map_items, self.unpack_items)

    def __call__(self, data):
        if len(self.transforms) == 0:
            return data
        index = self.R.multinomial(1, self.weights).argmax()
        _transform = self.transforms[index]
        data = apply_transform(_transform, data, self.map_items, self.unpack_items, self.log_stats)
        # if the data is a mapping (dictionary), append the OneOf transform to the end
        if isinstance(data, monai.data.MetaTensor):
            self.push_transform(data, extra_info={"index": index})
        elif isinstance(data, Mapping):
            for key in data:  # dictionary not change size during iteration
                if isinstance(data[key], monai.data.MetaTensor):
                    self.push_transform(data[key], extra_info={"index": index})
        return data

    def inverse(self, data):
        if len(self.transforms) == 0:
            return data

        index = None
        if isinstance(data, monai.data.MetaTensor):
            index = self.pop_transform(data)[TraceKeys.EXTRA_INFO]["index"]
        elif isinstance(data, Mapping):
            for key in data:
                if isinstance(data[key], monai.data.MetaTensor):
                    index = self.pop_transform(data, key)[TraceKeys.EXTRA_INFO]["index"]