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

Code before:
            for NumPy array and PyTorch Tensor, log the data shape and value range,
            for other metadata, log the values directly. Default to `False`.
        min_samples: minimum number of transforms to sample. Defaults to `0`.
        max_samples: maximum number of transforms to sample. Defaults to `3`.
        fixed: whether to sample exactly `max_samples` transforms, or up to it. Defaults to `False`.
    """

    def __init__(
        self,
        transforms: Sequence[Callable] | Callable | None = None,
        map_items: bool = True,
        unpack_items: bool = False,
        log_stats: bool = False,
        min_samples: int = 0,
        max_samples: int = 3,
        fixed: bool = True,
    ) -> None:
        super().__init__(transforms, map_items, unpack_items, log_stats)
        if transforms is None:
            transforms = []
        self.original_transforms = transforms
        self.n_transforms = len(transforms)
        self.max_samples = min(self.n_transforms, max_samples)
        self.fixed = fixed
        self.min_samples = min_samples
        self.random = random.Random(get_seed())

    def __call__(self, data):
        if self.n_transforms != 0:
            sample_size = self.max_samples if self.fixed else self.random.randint(self.min_samples, self.max_samples)
            self.active = self.random.sample(np.arange(self.n_transforms).tolist(), sample_size)
            self.transforms = [self.original_transforms[i] for i in self.active]

        result = super().__call__(data)
        return result

    def set_random_state(self, seed: int | None = None, state: np.random.RandomState | None = None) -> Compose:
        self.random = random.Random(seed)
        return super().set_random_state(seed, state)


Review comment (suggestion):
I would initialise `self.active` here so that it always exists rather than popping in when `__call__` is called for the first time. It's initial value can just be `np.arange(self.n_transforms)`.

Code after:
        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"]
        else:
            raise RuntimeError(
                f"Inverse only implemented for Mapping (dictionary) or MetaTensor data, got type {type(data)}."
            )
        if index is None:
            # no invertible transforms have been applied
            return data

        _transform = self.transforms[index]
        # apply the inverse