Harmonic

LinearHarmonicPrior

基类：PriorDistribution

一个子类，表示来自 Jit 等人的线性谐波先验分布。https://arxiv.org/abs/2304.02198。

源代码位于 bionemo/moco/distributions/prior/continuous/harmonic.py

class LinearHarmonicPrior(PriorDistribution):
    """A subclass representing a Linear Harmonic prior distribution from Jit et al. https://arxiv.org/abs/2304.02198."""

    def __init__(
        self,
        distance: Float = 3.8,
        length: Optional[int] = None,
        center: Bool = False,
        rng_generator: Optional[torch.Generator] = None,
        device: Union[str, torch.device] = "cpu",
    ) -> None:
        """Linear Harmonic prior distribution.

        Args:
            distance (Float): RMS distance between adjacent points in the line graph.
            length (Optional[int]): The number of points in a batch.
            center (bool): Whether to center the samples around the mean. Defaults to False.
            rng_generator: An optional :class:`torch.Generator` for reproducible sampling. Defaults to None.
            device (Optional[str]): Device to place the schedule on (default is "cpu").
        """
        self.distance = distance
        self.length = length
        self.center = center
        self.rng_generator = rng_generator
        self.device = device
        if length:
            self._calculate_terms(length, device)

    def _calculate_terms(self, N, device):
        a = 3 / (self.distance * self.distance)
        J = torch.zeros(N, N)
        for i, j in zip(torch.arange(N - 1), torch.arange(1, N)):
            J[i, i] += a
            J[j, j] += a
            J[i, j] = J[j, i] = -a
        D, P = torch.linalg.eigh(J)
        D_inv = 1 / D
        D_inv[0] = 0
        self.P, self.D_inv = P.to(device), D_inv.to(device)

    def sample(
        self,
        shape: Tuple,
        mask: Optional[Tensor] = None,
        device: Union[str, torch.device] = "cpu",
        rng_generator: Optional[torch.Generator] = None,
    ) -> Tensor:
        """Generates a specified number of samples from the Harmonic prior distribution.

        Args:
            shape (Tuple): The shape of the samples to generate.
            device (str): cpu or gpu.
            mask (Optional[Tensor]): An optional mask to apply to the samples. Defaults to None.
            rng_generator: An optional :class:`torch.Generator` for reproducible sampling. Defaults to None.

        Returns:
            Float: A tensor of samples.
        """
        if len(shape) != 3:
            raise ValueError("Input shape can only work for B x L x D")
        if rng_generator is None:
            rng_generator = self.rng_generator

        samples = torch.randn(*shape, device=device, generator=rng_generator)
        N = shape[1]

        if N != self.length:
            self._calculate_terms(N, device)

        std = torch.sqrt(self.D_inv).unsqueeze(-1)
        samples = self.P @ (std * samples)

        if self.center:
            samples = remove_center_of_mass(samples, mask)

        if mask is not None:
            samples = samples * mask.unsqueeze(-1)
        return samples

__init__(distance=3.8, length=None, center=False, rng_generator=None, device='cpu')

线性谐波先验分布。

参数

名称	类型	描述	默认值
distance	浮点	线图中相邻点之间的 RMS 距离。	3.8
length	Optional[int]	批次中的点数。	无
center	布尔值	是否将样本围绕均值居中。默认为 False。	False
rng_generator	Optional[Generator]	用于可重现采样的可选 :class:torch.Generator。默认为 None。	无
device	Optional[str]	放置调度的设备（默认为“cpu”）。	'cpu'

源代码位于 bionemo/moco/distributions/prior/continuous/harmonic.py

def __init__(
    self,
    distance: Float = 3.8,
    length: Optional[int] = None,
    center: Bool = False,
    rng_generator: Optional[torch.Generator] = None,
    device: Union[str, torch.device] = "cpu",
) -> None:
    """Linear Harmonic prior distribution.

    Args:
        distance (Float): RMS distance between adjacent points in the line graph.
        length (Optional[int]): The number of points in a batch.
        center (bool): Whether to center the samples around the mean. Defaults to False.
        rng_generator: An optional :class:`torch.Generator` for reproducible sampling. Defaults to None.
        device (Optional[str]): Device to place the schedule on (default is "cpu").
    """
    self.distance = distance
    self.length = length
    self.center = center
    self.rng_generator = rng_generator
    self.device = device
    if length:
        self._calculate_terms(length, device)

sample(shape, mask=None, device='cpu', rng_generator=None)

从 Harmonic 先验分布生成指定数量的样本。

参数

名称	类型	描述	默认值
shape	元组	要生成的样本的形状。	必需
device	字符串	cpu 或 gpu。	'cpu'
mask	Optional[Tensor]	应用于样本的可选掩码。默认为 None。	无
rng_generator	Optional[Generator]	用于可重现采样的可选 :class:torch.Generator。默认为 None。	无

返回值

名称	类型	描述
浮点	张量	样本张量。

源代码位于 bionemo/moco/distributions/prior/continuous/harmonic.py

def sample(
    self,
    shape: Tuple,
    mask: Optional[Tensor] = None,
    device: Union[str, torch.device] = "cpu",
    rng_generator: Optional[torch.Generator] = None,
) -> Tensor:
    """Generates a specified number of samples from the Harmonic prior distribution.

    Args:
        shape (Tuple): The shape of the samples to generate.
        device (str): cpu or gpu.
        mask (Optional[Tensor]): An optional mask to apply to the samples. Defaults to None.
        rng_generator: An optional :class:`torch.Generator` for reproducible sampling. Defaults to None.

    Returns:
        Float: A tensor of samples.
    """
    if len(shape) != 3:
        raise ValueError("Input shape can only work for B x L x D")
    if rng_generator is None:
        rng_generator = self.rng_generator

    samples = torch.randn(*shape, device=device, generator=rng_generator)
    N = shape[1]

    if N != self.length:
        self._calculate_terms(N, device)

    std = torch.sqrt(self.D_inv).unsqueeze(-1)
    samples = self.P @ (std * samples)

    if self.center:
        samples = remove_center_of_mass(samples, mask)

    if mask is not None:
        samples = samples * mask.unsqueeze(-1)
    return samples