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
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
|
/*
* Licensed to the Apache Software Foundation (ASF) under one or more
* contributor license agreements. See the NOTICE file distributed with
* this work for additional information regarding copyright ownership.
* The ASF licenses this file to You under the Apache License, Version 2.0
* (the "License"); you may not use this file except in compliance with
* the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.apache.spark.util.random
import java.util.Random
import cern.jet.random.Poisson
import cern.jet.random.engine.DRand
/**
* A pseudorandom sampler. It is possible to change the sampled item type. For example, we might
* want to add weights for stratified sampling or importance sampling. Should only use
* transformations that are tied to the sampler and cannot be applied after sampling.
*
* @tparam T item type
* @tparam U sampled item type
*/
trait RandomSampler[T, U] extends Pseudorandom with Cloneable with Serializable {
/** take a random sample */
def sample(items: Iterator[T]): Iterator[U]
override def clone: RandomSampler[T, U] =
throw new NotImplementedError("clone() is not implemented.")
}
/**
* A sampler based on Bernoulli trials.
*
* @param lb lower bound of the acceptance range
* @param ub upper bound of the acceptance range
* @param complement whether to use the complement of the range specified, default to false
* @tparam T item type
*/
class BernoulliSampler[T](lb: Double, ub: Double, complement: Boolean = false)
(implicit random: Random = new XORShiftRandom)
extends RandomSampler[T, T] {
def this(ratio: Double)(implicit random: Random = new XORShiftRandom)
= this(0.0d, ratio)(random)
override def setSeed(seed: Long) = random.setSeed(seed)
override def sample(items: Iterator[T]): Iterator[T] = {
items.filter { item =>
val x = random.nextDouble()
(x >= lb && x < ub) ^ complement
}
}
override def clone = new BernoulliSampler[T](lb, ub)
}
/**
* A sampler based on values drawn from Poisson distribution.
*
* @param poisson a Poisson random number generator
* @tparam T item type
*/
class PoissonSampler[T](mean: Double)
(implicit var poisson: Poisson = new Poisson(mean, new DRand))
extends RandomSampler[T, T] {
override def setSeed(seed: Long) {
poisson = new Poisson(mean, new DRand(seed.toInt))
}
override def sample(items: Iterator[T]): Iterator[T] = {
items.flatMap { item =>
val count = poisson.nextInt()
if (count == 0) {
Iterator.empty
} else {
Iterator.fill(count)(item)
}
}
}
override def clone = new PoissonSampler[T](mean)
}
|
