PyTorch中自定义平面有界区域均匀分布的log_prob取值问题

针对平面上的有界区域均匀分布（比如你提到的下三角形区域：0≤x≤1、0≤y≤1且x≤y），自定义torch.distributions.Distribution类的log_prob方法可以遵循以下规则，和盒型均匀分布的逻辑对齐：

• 采样点不在目标区域内：直接返回-inf，和盒型分布处理越界点的方式一致。
• 采样点在目标区域内：返回该区域面积的负对数，也就是-log(区域面积)。

以你说的下三角形为例：
这个区域是单位正方形的一半，面积为0.5，所以符合条件的点的log_prob值为-log(0.5)（计算后约为0.6931）；任何不满足0≤x≤1、0≤y≤1或x≤y的点，都返回-inf。

下面是对应自定义分布的代码示例：

import torch
from torch.distributions import Distribution

class TriangleUniform(Distribution):
    def __init__(self):
        super().__init__(batch_shape=torch.Size(), event_shape=torch.Size([2]))
        # 预计算区域面积的对数
        self.log_area = torch.log(torch.tensor(0.5))

    def log_prob(self, value):
        x = value[..., 0]
        y = value[..., 1]
        # 判定采样点是否在目标区域内
        in_region = (x >= 0) & (x <= 1) & (y >= 0) & (y <= 1) & (x <= y)
        # 初始化结果为-inf，再给符合条件的点赋值
        log_p = torch.full_like(x, -float('inf'))
        log_p[in_region] = -self.log_area
        return log_p

# 测试用例
dist = TriangleUniform()
print(dist.log_prob(torch.tensor([0.3, 0.7])))  # 符合条件，输出tensor(0.6931)
print(dist.log_prob(torch.tensor([0.8, 0.4])))  # x>y，输出tensor(-inf)
print(dist.log_prob(torch.tensor([-0.2, 0.5]))) # x<0，输出tensor(-inf)

通用扩展 #

对于任意有界的高维空间均匀分布，逻辑都是一样的：

1. 先计算目标区域的体积（平面是面积，三维是体积，以此类推）。
2. 对每个采样点，判断是否满足区域的所有约束条件。
3. 满足条件的返回-log(区域体积)，不满足的返回-inf。

内容的提问来源于stack exchange，提问作者Erithacus Rubecula