Beta

BetaTimeDistribution

基类：TimeDistribution

表示 Beta 时间分布的类。

源代码在 bionemo/moco/distributions/time/beta.py 中

class BetaTimeDistribution(TimeDistribution):
    """A class representing a beta time distribution."""

    def __init__(
        self,
        p1: Float = 2.0,
        p2: Float = 1.0,
        min_t: Float = 0.0,
        max_t: Float = 1.0,
        discrete_time: Bool = False,
        nsteps: Optional[int] = None,
        rng_generator: Optional[torch.Generator] = None,
    ):
        """Initializes a BetaTimeDistribution object.

        Args:
            p1 (Float): The first shape parameter of the beta distribution.
            p2 (Float): The second shape parameter of the beta distribution.
            min_t (Float): The minimum time value.
            max_t (Float): The maximum time value.
            discrete_time (Bool): Whether the time is discrete.
            nsteps (Optional[int]): Number of nsteps for discretization.
            rng_generator: An optional :class:`torch.Generator` for reproducible sampling. Defaults to None.
        """
        super().__init__(discrete_time, nsteps, min_t, max_t, rng_generator)
        self.dist = torch.distributions.Beta(p1, p2)

    def sample(
        self, n_samples: int, device: Union[str, torch.device] = "cpu", rng_generator: Optional[torch.Generator] = None
    ):
        """Generates a specified number of samples from the uniform time distribution.

        Args:
            n_samples (int): The number of samples to generate.
            device (str): cpu or gpu.
            rng_generator: An optional :class:`torch.Generator` for reproducible sampling. Defaults to None.

        Returns:
            A tensor of samples.
        """
        if rng_generator is None:
            rng_generator = self.rng_generator
        time_step = self.dist.sample(torch.Size([n_samples])).to(device=device)
        if self.min_t and self.max_t and self.min_t > 0:
            time_step = time_step * (self.max_t - self.min_t) + self.min_t
        if self.discrete_time:
            if self.nsteps is None:
                raise ValueError("nsteps cannot be None for discrete time sampling")
            time_step = float_time_to_index(time_step, self.nsteps)
        return time_step

__init__(p1=2.0, p2=1.0, min_t=0.0, max_t=1.0, discrete_time=False, nsteps=None, rng_generator=None)

初始化 BetaTimeDistribution 对象。

参数

名称	类型	描述	默认值
p1	浮点数	beta 分布的第一个形状参数。	2.0
p2	浮点数	beta 分布的第二个形状参数。	1.0
min_t	浮点数	最小时间值。	0.0
max_t	浮点数	最大时间值。	1.0
discrete_time	布尔值	时间是否是离散的。	False
nsteps	Optional[int]	离散化的 nsteps 数量（可选）。	None
rng_generator	Optional[Generator]	用于可重现采样的可选 :class:torch.Generator。默认为 None。	None

源代码在 bionemo/moco/distributions/time/beta.py 中

def __init__(
    self,
    p1: Float = 2.0,
    p2: Float = 1.0,
    min_t: Float = 0.0,
    max_t: Float = 1.0,
    discrete_time: Bool = False,
    nsteps: Optional[int] = None,
    rng_generator: Optional[torch.Generator] = None,
):
    """Initializes a BetaTimeDistribution object.

    Args:
        p1 (Float): The first shape parameter of the beta distribution.
        p2 (Float): The second shape parameter of the beta distribution.
        min_t (Float): The minimum time value.
        max_t (Float): The maximum time value.
        discrete_time (Bool): Whether the time is discrete.
        nsteps (Optional[int]): Number of nsteps for discretization.
        rng_generator: An optional :class:`torch.Generator` for reproducible sampling. Defaults to None.
    """
    super().__init__(discrete_time, nsteps, min_t, max_t, rng_generator)
    self.dist = torch.distributions.Beta(p1, p2)

sample(n_samples, device='cpu', rng_generator=None)

从均匀时间分布生成指定数量的样本。

参数

名称	类型	描述	默认值
n_samples	整数	要生成的样本数。	必需
device	字符串	cpu 或 gpu。	'cpu'
rng_generator	Optional[Generator]	用于可重现采样的可选 :class:torch.Generator。默认为 None。	None

返回值

类型	描述
	样本张量。

源代码在 bionemo/moco/distributions/time/beta.py 中

def sample(
    self, n_samples: int, device: Union[str, torch.device] = "cpu", rng_generator: Optional[torch.Generator] = None
):
    """Generates a specified number of samples from the uniform time distribution.

    Args:
        n_samples (int): The number of samples to generate.
        device (str): cpu or gpu.
        rng_generator: An optional :class:`torch.Generator` for reproducible sampling. Defaults to None.

    Returns:
        A tensor of samples.
    """
    if rng_generator is None:
        rng_generator = self.rng_generator
    time_step = self.dist.sample(torch.Size([n_samples])).to(device=device)
    if self.min_t and self.max_t and self.min_t > 0:
        time_step = time_step * (self.max_t - self.min_t) + self.min_t
    if self.discrete_time:
        if self.nsteps is None:
            raise ValueError("nsteps cannot be None for discrete time sampling")
        time_step = float_time_to_index(time_step, self.nsteps)
    return time_step