pocketpy/src/random.cpp
2024-02-15 00:13:34 +08:00

101 lines
4.1 KiB
C++

#include "pocketpy/random.h"
namespace pkpy{
struct Random{
PY_CLASS(Random, random, Random)
std::mt19937 gen;
Random(){
gen.seed(std::chrono::high_resolution_clock::now().time_since_epoch().count());
}
static void _register(VM* vm, PyObject* mod, PyObject* type){
vm->bind_default_constructor<Random>(type);
vm->bind_method<1>(type, "seed", [](VM* vm, ArgsView args) {
Random& self = _CAST(Random&, args[0]);
self.gen.seed(CAST(i64, args[1]));
return vm->None;
});
vm->bind_method<2>(type, "randint", [](VM* vm, ArgsView args) {
Random& self = _CAST(Random&, args[0]);
i64 a = CAST(i64, args[1]);
i64 b = CAST(i64, args[2]);
if (a > b) vm->ValueError("randint(a, b): a must be less than or equal to b");
std::uniform_int_distribution<i64> dis(a, b);
return VAR(dis(self.gen));
});
vm->bind_method<0>(type, "random", [](VM* vm, ArgsView args) {
Random& self = _CAST(Random&, args[0]);
std::uniform_real_distribution<f64> dis(0.0, 1.0);
return VAR(dis(self.gen));
});
vm->bind_method<2>(type, "uniform", [](VM* vm, ArgsView args) {
Random& self = _CAST(Random&, args[0]);
f64 a = CAST(f64, args[1]);
f64 b = CAST(f64, args[2]);
std::uniform_real_distribution<f64> dis(a, b);
return VAR(dis(self.gen));
});
vm->bind_method<1>(type, "shuffle", [](VM* vm, ArgsView args) {
Random& self = _CAST(Random&, args[0]);
List& L = CAST(List&, args[1]);
std::shuffle(L.begin(), L.end(), self.gen);
return vm->None;
});
vm->bind_method<1>(type, "choice", [](VM* vm, ArgsView args) {
Random& self = _CAST(Random&, args[0]);
auto [data, size] = vm->_cast_array(args[1]);
if(size == 0) vm->IndexError("cannot choose from an empty sequence");
std::uniform_int_distribution<i64> dis(0, size - 1);
return data[dis(self.gen)];
});
vm->bind(type, "choices(self, population, weights=None, k=1)", [](VM* vm, ArgsView args) {
Random& self = _CAST(Random&, args[0]);
auto [data, size] = vm->_cast_array(args[1]);
if(size == 0) vm->IndexError("cannot choose from an empty sequence");
std::vector<f64> cum_weights(size);
if(args[2] == vm->None){
for(int i = 0; i < size; i++) cum_weights[i] = i + 1;
}else{
auto [weights, weights_size] = vm->_cast_array(args[2]);
if(weights_size != size) vm->ValueError(_S("len(weights) != ", size));
cum_weights[0] = CAST(f64, weights[0]);
for(int i = 1; i < size; i++){
cum_weights[i] = cum_weights[i - 1] + CAST(f64, weights[i]);
}
}
if(cum_weights[size - 1] <= 0) vm->ValueError("total of weights must be greater than zero");
int k = CAST(i64, args[3]);
List result(k);
for(int i = 0; i < k; i++){
f64 r = std::uniform_real_distribution<f64>(0.0, cum_weights[size - 1])(self.gen);
int idx = std::lower_bound(cum_weights.begin(), cum_weights.end(), r) - cum_weights.begin();
result[i] = data[idx];
}
return VAR(std::move(result));
});
}
};
void add_module_random(VM* vm){
PyObject* mod = vm->new_module("random");
Random::register_class(vm, mod);
PyObject* instance = vm->heap.gcnew<Random>(Random::_type(vm));
mod->attr().set("seed", vm->getattr(instance, "seed"));
mod->attr().set("random", vm->getattr(instance, "random"));
mod->attr().set("uniform", vm->getattr(instance, "uniform"));
mod->attr().set("randint", vm->getattr(instance, "randint"));
mod->attr().set("shuffle", vm->getattr(instance, "shuffle"));
mod->attr().set("choice", vm->getattr(instance, "choice"));
mod->attr().set("choices", vm->getattr(instance, "choices"));
}
} // namespace pkpy