summaryrefslogtreecommitdiffstats
path: root/eval/src/tests/eval/gbdt/model.cpp
blob: 112f058fa2c02bc74687d7da7741eb1565ae0534 (plain) (blame)
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
96
97
98
99
// Copyright 2017 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
#pragma once

#include <random>
#include <vespa/eval/eval/function.h>

using vespalib::make_string;
using vespalib::eval::Function;

//-----------------------------------------------------------------------------

class Model
{
private:
    std::mt19937 _gen;
    size_t _less_percent;

    size_t get_int(size_t min, size_t max) {
        std::uniform_int_distribution<size_t> dist(min, max);
        return dist(_gen);
    }

    double get_real(double min, double max) {
        std::uniform_real_distribution<double> dist(min, max);
        return dist(_gen);
    }

    std::string make_feature_name() {
        size_t max_feature = 2;
        while ((max_feature < 1024) && (get_int(0, 99) < 55)) {
            max_feature *= 2;
        }
        return make_string("feature_%zu", get_int(1, max_feature));
    }

    std::string make_cond() {
        if (get_int(1,100) > _less_percent) {
            return make_string("(%s in [%g,%g,%g])",
                               make_feature_name().c_str(),
                               get_int(0, 4) / 4.0,
                               get_int(0, 4) / 4.0,
                               get_int(0, 4) / 4.0);
        } else {
            return make_string("(%s<%g)",
                               make_feature_name().c_str(),
                               get_real(0.0, 1.0));
        }
    }

public:
    explicit Model(size_t seed = 5489u) : _gen(seed), _less_percent(80) {}

    Model &less_percent(size_t value) {
        _less_percent = value;
        return *this;
    }

    std::string make_tree(size_t size) {
        assert(size > 0);
        if (size == 1) {
            return make_string("%g", get_real(0.0, 1.0));
        }
        size_t pivot = get_int(1, size - 1);
        return make_string("if(%s,%s,%s)",
                           make_cond().c_str(),
                           make_tree(pivot).c_str(),
                           make_tree(size - pivot).c_str());
    }

    std::string make_forest(size_t num_trees, size_t tree_sizes) {
        assert(num_trees > 0);
        vespalib::string forest = make_tree(tree_sizes);
        for (size_t i = 1; i < num_trees; ++i) {
            forest.append("+");
            forest.append(make_tree(tree_sizes));
        }
        return forest;
    }
};

//-----------------------------------------------------------------------------

struct ForestParams {
    size_t model_seed;
    size_t less_percent;
    size_t tree_size;
    ForestParams(size_t model_seed_in, size_t less_percent_in, size_t tree_size_in)
        : model_seed(model_seed_in), less_percent(less_percent_in), tree_size(tree_size_in) {}
};

//-----------------------------------------------------------------------------

Function make_forest(const ForestParams &params, size_t num_trees) {
    return Function::parse(Model(params.model_seed)
                           .less_percent(params.less_percent)
                           .make_forest(num_trees, params.tree_size));
}

//-----------------------------------------------------------------------------