bed-reader 1.0.6

Read and write the PLINK BED format, simply and efficiently.
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
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
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135

import time
from pathlib import Path

import numpy as np
import pandas as pd

from bed_reader import create_bed, to_bed

if True:
    ssd_path = Path("m:/deldir/bench")
    hdd_path = Path("e:/deldir/bench")
else:
    ssd_path = Path("/mnt/m/deldir/bench")
    hdd_path = Path("/mnt/e/deldir/bench")


def test_writes(iid_count, sid_count, num_threads, stream, drive, include_error):
    if drive == "ssd":
        path = ssd_path
    elif drive == "hdd":
        path = hdd_path
    else:
        raise ValueError(f"drive must be 'ssd' or 'hdd', not '{drive}'")

    output_file = path / f"{iid_count}x{sid_count}.bed"

    val = np.full((iid_count, sid_count), 1, dtype=np.int8, order="F")
    if include_error:
        val[iid_count // 2, sid_count // 2] = 22

    val_size = float(iid_count) * sid_count
    if val_size > 9_200_000_000:
        raise ValueError(f"val_size {val_size} is too large")

    result_list = []
    start = time.time()
    if not stream:
        to_bed(output_file, val, num_threads=num_threads)
    else:
        with create_bed(
            output_file,
            iid_count=iid_count,
            sid_count=sid_count,
            num_threads=num_threads,
        ) as bed_writer:
            for column_data in val.T:
                bed_writer.write(column_data)
    delta = time.time() - start
    result = [
        iid_count,
        sid_count,
        num_threads,
        stream,
        drive,
        include_error,
        True,
        "i8",
        "benchmark.py",
        val_size,
        round(delta, 4),
        round(val_size / delta),
        round(val_size / delta / num_threads),
    ]
    print(result)
    result_list.append(result)

    result_df = pd.DataFrame(
        result_list,
        columns=[
            "iid_count",
            "sid_count",
            "num_threads",
            "stream",
            "drive",
            "include_error",
            "release",
            "dtype",
            "source",
            "val size",
            "time",
            "val per second",
            "per thread",
        ],
    )
    return result_df


def meta_test(
    iid_count,
    sid_start=5,
    sid_end=None,
    point_count=5,  # 30
    drive_list=["ssd"],
    plot_index=0,
):
    # 5K vs 50K
    # 50K vs 50K
    # 500K vs 5K
    if sid_end is None:
        sid_end = 50000 if iid_count <= 50000 else 5000

    result = []
    for sid_count in np.logspace(np.log10(sid_start), np.log10(sid_end), point_count, base=10, dtype=int):
        for drive in drive_list:
            for num_threads in [1, 12]:
                for stream in [True, False]:
                    result.append(test_writes(iid_count, sid_count, num_threads, stream, drive, False))
    df = pd.concat(result)
    df.to_csv(ssd_path / "plots" / f"bench{plot_index},iid_count{iid_count}.csv", index=False)
    df2 = df.pivot(
        index="sid_count",
        columns=["iid_count", "drive", "num_threads", "stream"],
        values="val per second",
    )
    df2.plot(marker=".", logx=True)
    plt.savefig(ssd_path / "plots" / f"plot{plot_index},iid_count{iid_count},{'_'.join(drive_list)}.png")
    # plt.show()
    return df


if __name__ == "__main__":
    import matplotlib
    import matplotlib.pyplot as plt

    matplotlib.use("Agg")

    plot_count = 0
    for drive in ["ssd"]:  # , "hdd"
        for iid_count in [50_000]:  # 5_000, 500_000]:
            meta_test(
                iid_count,
                drive_list=[drive],
                plot_index=plot_count,
            )
            plot_count += 1