Source code for bayesflow.adapters.transforms.random_subsample

import keras
import keras.ops as ops
from bayesflow.utils.serialization import serializable, serialize
from bayesflow.types import Tensor

from .elementwise_transform import ElementwiseTransform




@serializable(package="bayesflow.adapters")
class RandomSubsample(ElementwiseTransform):
    """
    A transform that takes a random subsample of the data within an axis.

    Examples
    --------

    >>> adapter = bf.Adapter().random_subsample("x", sample_size=3, axis=-1)
    """

    def __init__(
        self,
        sample_size: int | float,
        axis: int = -1,
    ):
        super().__init__()
        if isinstance(sample_size, float):
            if sample_size <= 0 or sample_size >= 1:
                raise ValueError("Sample size as a percentage must be a float between 0 and 1 exclusive. ")
        self.sample_size = sample_size
        self.axis = axis



    def forward(self, data: Tensor, **kwargs) -> Tensor:
        max_sample_size = ops.shape(data)[self.axis]

        if isinstance(self.sample_size, int):
            sample_size = self.sample_size
        else:
            sample_size = int(ops.round(self.sample_size * max_sample_size))

        # random sample without replacement
        indices = keras.random.shuffle(ops.arange(max_sample_size))[:sample_size]
        return ops.take(data, indices, axis=self.axis)




    def inverse(self, data: Tensor, **kwargs) -> Tensor:
        # non invertible transform
        return data




    def get_config(self) -> dict:
        config = {"sample_size": self.sample_size, "axis": self.axis}

        return serialize(config)