正好我上周刚写了一个同样的算法, 是对这种随机抽样的进一步优化.
对于一个任意形式的 PDF 函数, 如果直接使用 $g(x) = Uniform(x_ {min}, x_{max})$ 作为采样取样, 则命中率过低, 效率低下 (尤其是存在较长拖尾的函数).
因此, 我们可以构建一个分段函数 g(x), 每一段都是一个均匀分布的函数, 它形式简单, 易于生成和抽样.
如上图中红色曲线所示的, 灰色阴影(及蓝色区域) 围成的面积, 就是我们的 g(x) , 它始终有 g(x) < f(x) . 下半部分所示的就是 g(x) 的 CDF 曲线, 我们只需要记录折线的每一个端点 (x_i, cdf_i) . 对于任意服从均匀分布的变量 p, 我们找到它所在的区间 $(x_ i, cdf_ i) \sim (x_ {i+1}, cdf_ {i+1})$, 也就可以计算出对应的 $x = x_ i + \frac{x_ {i+1} - x_ i}{cdf_ {i+1} - cdf_ i} (p-cdf_ i)$ . 这样得到的 x 也就是服从分布 p(x) 的随机变量.
根据 $x \sim p(x) $ 取样得到 f(x) 的过程类似: 取得随机变量 $y \sim Uniform(0, g(x))$ , 若 $y > f(x)$ 则舍弃该组取样, 重新取一个 x 样本 (图中 y_drop) , 反之, 当 $y < f(x)$ 时 , 可以返回该随机变量 x. 如此一来, 我们最后获得的 x 就是服从 f(x) 分布的随机变量.
至于 p(x) 函数的取得, 考虑到程序精度等问题, 我采用了如下的方法.
- 对于每一段 [x_i, x_i+1] , 我们在该区间内随机取 50 个点, 再每隔 (dx / 50) 取一个点, 共 100 个点.
- 计算这 100 个点对应的 y_k=f(x_k) 值, 寻找最大值 y_max,
- 将 y_max 乘 1.1 , 记为 y_i, 保证 $y_i > f(x) , x \in [x_i, x_{i+1}] $
- 则每段 p(x) 使用的估计值为 y_i
如下是使用 Typescript 代码的实现, ./random.ts:
// import 使用的是 Deno 风格的, 如果是 npm 风格需要去掉 "npm:"
import * as d3rnd from "npm:d3-random";
import seedrandom from "npm:seedrandom";
// 设置种子（可选）
const rng = seedrandom("some-seed");
// 创建一个均匀分布采样器（范围 0 到 1）
const uniform = d3rnd.randomUniform.source(rng)(0, 1);
class ProbabilityDistribution {
f: (x: number) => number;
x_min: number;
x_max: number;
config: {
num_segment: number;
num_prob: number;
ratio_increase: number;
max_step: number;
};
private interval: number;
private segment_start: number[];
// g(x) of each segment
private segments: number[];
// accumulation of g(x) from first segment
private segments_acc: number[];
private sample_info = {
/** 执行取样次数 */
sample_count: 0,
/** 取样时尝试的次数 */
trail_count: 0,
/** 取样时失败的次数 (尝试次数超过最大步数限制) */
fail_count: 0,
};
constructor(f: (x: number) => number, x_min: number, x_max: number, {
num_segment = 10,
num_prob = 50,
ratio_increase = 1.1,
max_step = 100,
}) {
this.f = f;
this.x_min = x_min;
this.x_max = x_max;
this.config = { num_segment, num_prob, ratio_increase, max_step };
this.interval = (x_max - x_min) / num_segment;
this.segment_start = [];
this.segments = [];
this.segments_acc = [0];
this.build();
if (!this.check()) {
throw new Error("Failed to create Probability Distribution!");
}
}
private build() {
let total_y = 0;
for (let i = 0; i < this.config.num_segment; i++) {
const x = this.x_min + this.interval * i;
let y = Math.max(this.f(x), 0);
for (let j = 0; j < this.config.num_prob; j++) {
const y2 = this.f(x + this.interval * j / this.config.num_prob);
if (y2 > y) {
y = y2;
}
const y3 = this.f(x + this.interval * Math.random());
if (y3 > y) {
y = y3;
}
}
y *= this.config.ratio_increase;
total_y += y;
this.segment_start.push(x);
this.segments.push(y);
this.segments_acc.push(total_y);
}
for (let i = 0; i <= this.segments.length; i++) {
this.segments_acc[i] /= total_y;
}
}
check(): boolean {
function is_pos(x: number) {
return x >= 0 && Number.isFinite(x);
}
function is_good(x: number) {
return x >= 0 && x <= 1;
}
if (this.interval <= 0 || !Number.isFinite(this.interval)) {
console.error(`Error: interval is ${this.interval}`);
return false;
}
if (
this.x_min >= this.x_max || !Number.isFinite(this.x_min) ||
!Number.isFinite(this.x_max)
) {
console.log(`Error: range is [${this.x_min},${this.x_max}]`);
return false;
}
for (let i = 0; i < this.segments.length; i++) {
if (!is_pos(this.segments[i])) {
console.error(`Error: segment[${i}] is ${this.segments[i]}`);
return false;
}
}
for (let i = 0; i < this.segments_acc.length; i++) {
if (!is_good(this.segments_acc[i])) {
console.error(`Error: segment_acc[${i}] is ${this.segments_acc[i]}`);
return false;
}
}
return true;
}
/**
* Find the segment position of x
* @param x x in [0, 1]
* @returns [segment_index, pos_in_segment]. Index in [0, num_segment], Pos in [0, 1).
*/
findSegment(x: number): [number, number] {
if (x < 0 || x > 1) {
throw new Error("x must be in [0, 1]");
}
// Binary search for the rightmost arr[i] <= x
let low = 0, high = this.segments_acc.length;
while (low < high) {
const mid = Math.ceil((low + high) / 2);
if (this.segments_acc[mid] <= x) {
low = mid;
} else {
high = mid - 1;
}
}
// arr[low] <= x < arr[low + 1]
const a = this.segments_acc[low], b = this.segments_acc[low + 1];
return [low, (x - a) / (b - a)];
}
sample(): number | null {
this.sample_info.sample_count += 1;
let i = 0;
while (i++ < this.config.max_step) {
this.sample_info.trail_count += 1;
const [segment, pos] = this.findSegment(randomUniform());
const x = this.x_min + (segment + pos) * this.interval;
const fy = this.f(x);
const y = randomUniform(0, this.segments[segment]);
if (y < fy) {
return x;
}
}
this.sample_info.fail_count += 1;
return null;
}
info() {
const accept_rate =
(this.sample_info.sample_count - this.sample_info.fail_count) /
this.sample_info.trail_count;
const estimated_fail = Math.pow(
1 / (1 - accept_rate),
this.config.max_step,
);
const info = {
accept_rate,
...this.sample_info,
estimated_fail,
};
return info;
}
}
export { ProbabilityDistribution };
测试代码如下:
function test(){
const config = { num_segment: 10, num_prob: 50 };
const dist = new ProbabilityDistribution(
(x) => 1 / (1 + x * x),
-1,
10,
config,
);
const start = performance.now();
const arr = Array.from({ length: 1e7 }, (_) => dist.sample());
const end = performance.now();
console.log(`Time to gen ${arr.length} samples: ${end - start} ms`);
console.log("Info:", dist.info());
}
测试结果如下
Time to gen 10000000 samples: 1454.5922 ms
Info: {
accept_rate: 0.6739182062777368,
sample_count: 10000000,
trail_count: 14838596,
fail_count: 0,
estimated_fail: 4.648842881822224e+48
}
可以通过提高 num_segment (分段数量) 的方式提高接受率, 进而降低时间成本.
Note: 上述代码中使用了 npm 库 seedrandom 和 d3-random , 这是我测试得到的生成随机数的质量较高的方式. 如你不介意质量, 也可以直接使用 Math.random() 提高效率.
