/
ztp.js
46 lines (43 loc) · 901 Bytes
/
ztp.js
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
/**
* Zero-truncated poisson
*/
export default class ZeroTruncatedPoisson {
// https://en.wikipedia.org/wiki/Zero-truncated_Poisson_distribution
// https://www.stat.umn.edu/geyer/3701/notes/zero.html
constructor() {
this._lr = 0.1
}
/**
* Fit model.
*
* @param {number[]} x Training data
*/
fit(x) {
const mean = x.reduce((s, v) => s + v, 0) / x.length
let th = 0
while (true) {
const m = Math.exp(th)
const myu = m / (1 - Math.exp(-m))
th += this._lr * (mean - myu)
if (Math.abs(mean - myu) < 1.0e-8) {
break
}
}
this._lambda = Math.exp(th)
}
/**
* Returns predicted probabilities.
*
* @param {number[]} x Sample data
* @returns {number[]} Predicted values
*/
probability(x) {
return x.map(v => {
let f = 1
for (let i = 2; i <= v; i++) {
f *= i
}
return this._lambda ** v / ((Math.exp(this._lambda) - 1) * f)
})
}
}