inferust 0.1.12

Statistical modeling for Rust — OLS/WLS regression, GLM, survival analysis, ARIMA/VAR, nonparametric tests, and more. A statsmodels-style library.
Documentation 
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

#!/usr/bin/env python3
import argparse
import time

import numpy as np
import statsmodels.api as sm


def dataset(n: int, p: int):
    i = np.arange(n, dtype=np.float64)[:, None]
    j = np.arange(p, dtype=np.float64)[None, :]
    x = np.sin((i + 1.0) * (j + 1.0) * 0.001)
    x += np.cos((i + j + 3.0) * 0.017)
    x += (np.mod(i, 97.0)) * 0.0001
    beta = (np.arange(p, dtype=np.float64) + 1.0) / float(p)
    y = 1.5 + x @ beta + np.sin((np.arange(n, dtype=np.float64) + 11.0) * 0.037) * 0.01
    return x, y


def millis(seconds: float) -> float:
    return seconds * 1000.0


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--rows", type=int, default=10_000)
    parser.add_argument("--features", type=int, default=8)
    parser.add_argument("--repeats", type=int, default=10)
    parser.add_argument("--warmups", type=int, default=2)
    args = parser.parse_args()

    x, y = dataset(args.rows, args.features)
    x_const = sm.add_constant(x, has_constant="add")

    for _ in range(args.warmups):
        sm.OLS(y, x_const).fit()

    timings = []
    checksum = 0.0
    for _ in range(args.repeats):
        started = time.perf_counter()
        result = sm.OLS(y, x_const).fit()
        elapsed = time.perf_counter() - started
        checksum += float(result.params.sum())
        timings.append(elapsed)

    timings.sort()
    print(
        "engine=python-statsmodels "
        f"rows={args.rows} features={args.features} repeats={args.repeats} warmups={args.warmups} "
        f"min_ms={millis(timings[0]):.3f} "
        f"median_ms={millis(timings[len(timings) // 2]):.3f} "
        f"mean_ms={millis(sum(timings) / len(timings)):.3f} "
        f"checksum={checksum:.8f}"
    )


if __name__ == "__main__":
    main()