You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
120 lines
4.1 KiB
120 lines
4.1 KiB
// For licensing see accompanying LICENSE.md file.
|
|
// Copyright (C) 2022 Apple Inc. All Rights Reserved.
|
|
|
|
import Foundation
|
|
import CoreML
|
|
|
|
/// A random source consistent with NumPy
|
|
///
|
|
/// This implementation matches:
|
|
/// [NumPy's older randomkit.c](https://github.com/numpy/numpy/blob/v1.0/numpy/random/mtrand/randomkit.c)
|
|
///
|
|
@available(iOS 16.2, macOS 13.1, *)
|
|
struct NumPyRandomSource: RandomNumberGenerator, RandomSource {
|
|
|
|
struct State {
|
|
var key = [UInt32](repeating: 0, count: 624)
|
|
var pos: Int = 0
|
|
var nextGauss: Double? = nil
|
|
}
|
|
|
|
var state: State
|
|
|
|
/// Initialize with a random seed
|
|
///
|
|
/// - Parameters
|
|
/// - seed: Seed for underlying Mersenne Twister 19937 generator
|
|
/// - Returns random source
|
|
init(seed: UInt32) {
|
|
state = .init()
|
|
var s = seed & 0xffffffff
|
|
for i in 0 ..< state.key.count {
|
|
state.key[i] = s
|
|
s = UInt32((UInt64(1812433253) * UInt64(s ^ (s >> 30)) + UInt64(i) + 1) & 0xffffffff)
|
|
}
|
|
state.pos = state.key.count
|
|
state.nextGauss = nil
|
|
}
|
|
|
|
/// Generate next UInt32 using fast 32bit Mersenne Twister
|
|
mutating func nextUInt32() -> UInt32 {
|
|
let n = 624
|
|
let m = 397
|
|
let matrixA: UInt64 = 0x9908b0df
|
|
let upperMask: UInt32 = 0x80000000
|
|
let lowerMask: UInt32 = 0x7fffffff
|
|
|
|
var y: UInt32
|
|
if state.pos == state.key.count {
|
|
for i in 0 ..< (n - m) {
|
|
y = (state.key[i] & upperMask) | (state.key[i + 1] & lowerMask)
|
|
state.key[i] = state.key[i + m] ^ (y >> 1) ^ UInt32((UInt64(~(y & 1)) + 1) & matrixA)
|
|
}
|
|
for i in (n - m) ..< (n - 1) {
|
|
y = (state.key[i] & upperMask) | (state.key[i + 1] & lowerMask)
|
|
state.key[i] = state.key[i + (m - n)] ^ (y >> 1) ^ UInt32((UInt64(~(y & 1)) + 1) & matrixA)
|
|
}
|
|
y = (state.key[n - 1] & upperMask) | (state.key[0] & lowerMask)
|
|
state.key[n - 1] = state.key[m - 1] ^ (y >> 1) ^ UInt32((UInt64(~(y & 1)) + 1) & matrixA)
|
|
state.pos = 0
|
|
}
|
|
y = state.key[state.pos]
|
|
state.pos += 1
|
|
|
|
y ^= (y >> 11)
|
|
y ^= (y << 7) & 0x9d2c5680
|
|
y ^= (y << 15) & 0xefc60000
|
|
y ^= (y >> 18)
|
|
|
|
return y
|
|
}
|
|
|
|
mutating func next() -> UInt64 {
|
|
let low = nextUInt32()
|
|
let high = nextUInt32()
|
|
return (UInt64(high) << 32) | UInt64(low)
|
|
}
|
|
|
|
/// Generate next random double value
|
|
mutating func nextDouble() -> Double {
|
|
let a = Double(nextUInt32() >> 5)
|
|
let b = Double(nextUInt32() >> 6)
|
|
return (a * 67108864.0 + b) / 9007199254740992.0
|
|
}
|
|
|
|
/// Generate next random value from a standard normal
|
|
mutating func nextGauss() -> Double {
|
|
if let nextGauss = state.nextGauss {
|
|
state.nextGauss = nil
|
|
return nextGauss
|
|
}
|
|
var x1, x2, r2: Double
|
|
repeat {
|
|
x1 = 2.0 * nextDouble() - 1.0
|
|
x2 = 2.0 * nextDouble() - 1.0
|
|
r2 = x1 * x1 + x2 * x2
|
|
} while r2 >= 1.0 || r2 == 0.0
|
|
|
|
// Box-Muller transform
|
|
let f = sqrt(-2.0 * log(r2) / r2)
|
|
state.nextGauss = f * x1
|
|
return f * x2
|
|
}
|
|
|
|
/// Generates a random value from a normal distribution with given mean and standard deviation.
|
|
mutating func nextNormal(mean: Double = 0.0, stdev: Double = 1.0) -> Double {
|
|
nextGauss() * stdev + mean
|
|
}
|
|
|
|
/// Generates an array of random values from a normal distribution with given mean and standard deviation.
|
|
mutating func normalArray(count: Int, mean: Double = 0.0, stdev: Double = 1.0) -> [Double] {
|
|
(0 ..< count).map { _ in nextNormal(mean: mean, stdev: stdev) }
|
|
}
|
|
|
|
/// Generate a shaped array with scalars from a normal distribution with given mean and standard deviation.
|
|
mutating func normalShapedArray(_ shape: [Int], mean: Double = 0.0, stdev: Double = 1.0) -> MLShapedArray<Double> {
|
|
let count = shape.reduce(1, *)
|
|
return .init(scalars: normalArray(count: count, mean: mean, stdev: stdev), shape: shape)
|
|
}
|
|
}
|