tetration 0.1.2

Tetration tensor file format: Rust library (tetration) and tet CLI
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
136
137
138
139
140
141
142
143
144
145
146

"""Native source reads for benchmark comparisons."""

from __future__ import annotations

import math
from dataclasses import dataclass

import numpy as np

from benchmark.case import BenchCase
from benchmark.constants import OpName, op_compat
from benchmark.log import time_warm_logged
from generate.constants import CHUNK_ELEMS
from generate.util import iter_slabs


@dataclass
class RunningStats:
    count: int = 0
    sum: float = 0.0
    min: float = math.inf
    max: float = -math.inf
    m2: float = 0.0

    def push_slab(self, arr: np.ndarray) -> None:
        flat = np.asarray(arr, dtype=np.float64).ravel()
        n = int(flat.size)
        if n == 0:
            return
        slab_sum = float(flat.sum())
        slab_mean = slab_sum / n
        slab_m2 = float(((flat - slab_mean) ** 2).sum())
        slab_min = float(flat.min())
        slab_max = float(flat.max())

        if self.count == 0:
            self.count = n
            self.sum = slab_sum
            self.min = slab_min
            self.max = slab_max
            self.m2 = slab_m2
            return

        n_a = self.count
        n_b = n
        total = n_a + n_b
        mean_a = self.sum / n_a
        delta = slab_mean - mean_a
        self.m2 += slab_m2 + delta * delta * n_a * n_b / total
        self.count = total
        self.sum += slab_sum
        self.min = min(self.min, slab_min)
        self.max = max(self.max, slab_max)

    def finish(self, op: OpName) -> float:
        if self.count == 0:
            raise ValueError("empty selection")
        match op:
            case "mean":
                return self.sum / self.count
            case "sum":
                return self.sum
            case "min":
                return self.min
            case "max":
                return self.max
            case "count":
                return float(self.count)
            case "var":
                return self.m2 / self.count if self.count else 0.0
            case "std":
                return math.sqrt(self.m2 / self.count) if self.count else 0.0
        raise ValueError(op)


def reduce_source_op(case: BenchCase, op: OpName) -> tuple[float, float]:
    compat = op_compat(case.format, op)
    if not compat.bench_source:
        raise ValueError(compat.note or "source not supported")

    slab_mib = CHUNK_ELEMS // (1024 * 1024)

    def run_h5() -> tuple[float, float]:
        import h5py

        with h5py.File(case.src, "r") as f:
            d = f["data"]
            length = int(d.shape[0])

            def compute() -> float:
                stats = RunningStats()
                for start, end in iter_slabs(length):
                    stats.push_slab(d[start:end])
                return stats.finish(op)

            return time_warm_logged(
                f"source {op} ({slab_mib} MiB slabs, h5py)",
                compute,
            )

    def run_netcdf() -> tuple[float, float]:
        import netCDF4 as nc

        with nc.Dataset(case.src, "r") as ds:
            var = ds.variables["data"]
            length = int(var.shape[0])

            def compute() -> float:
                stats = RunningStats()
                for start, end in iter_slabs(length):
                    stats.push_slab(var[start:end])
                return stats.finish(op)

            return time_warm_logged(
                f"source {op} ({slab_mib} MiB slabs, netCDF4)",
                compute,
            )

    def run_zarr() -> tuple[float, float]:
        import zarr

        root = zarr.open_group(str(case.src), mode="r")
        arr = root["data"]
        length = int(arr.shape[0])
        chunk_elems = int(arr.chunks[0])
        chunk_mib = chunk_elems * 4 // (1024 * 1024)

        def compute() -> float:
            stats = RunningStats()
            for start in range(0, length, chunk_elems):
                end = min(start + chunk_elems, length)
                stats.push_slab(arr[start:end])
            return stats.finish(op)

        return time_warm_logged(
            f"source {op} ({chunk_mib} MiB zarr chunks, raw)",
            compute,
        )

    match case.format:
        case "h5":
            return run_h5()
        case "netcdf":
            return run_netcdf()
        case "zarr":
            return run_zarr()