tiny.simd.stats

1pub const stats = @import("stats.zig");

1const std = @import("std");23pub fn Bins(comptime capacity: usize) type {4    return TypedBins(capacity, u32);5}67pub fn TypedBins(comptime capacity: usize, comptime Count: type) type {8    comptime requireBins(capacity, Count);910    return struct {11        counts: [capacity]Count = @splat(0),1213        const Self = @This();14        pub const bin_count = capacity;15        pub const CountType = Count;1617        pub fn notify(self: *Self, bin_value: anytype) void {18            const bin_index = index(bin_value);19            self.counts[bin_index] +%= 1;20        }2122        pub fn incrementBy(self: *Self, bin_value: anytype, count: u32) void {23            const bin_index = index(bin_value);24            self.counts[bin_index] +%= countValue(count);25        }2627        pub fn bin(self: *const Self, bin_index: usize) Count {28            std.debug.assert(bin_index < capacity);29            return self.counts[bin_index];30        }3132        pub fn resetBin(self: *Self, bin_index: usize) void {33            std.debug.assert(bin_index < capacity);34            self.counts[bin_index] = 0;35        }3637        pub fn assimilate(self: *Self, other: *const Self) void {38            for (&self.counts, other.counts) |*count, other_count| {39                count.* +%= other_count;40            }41        }4243        pub fn firstNonzero(self: *const Self) usize {44            for (self.counts, 0..) |count, bin_index| {45                if (count != 0) return bin_index;46            }47            return capacity;48        }4950        pub fn lastNonzero(self: *const Self) usize {51            var bin_index = capacity;52            while (bin_index != 0) {53                bin_index -= 1;54                if (self.counts[bin_index] != 0) return bin_index;55            }56            return 0;57        }5859        pub fn numNonzero(self: *const Self) usize {60            var count: usize = 0;61            for (self.counts) |bin_count_value| count += @intFromBool(bin_count_value != 0);62            return count;63        }6465        pub fn modalBinIndex(self: *const Self) usize {66            var maximum: Count = 0;67            var maximum_index: usize = 0;68            for (self.counts, 0..) |count, bin_index| {69                if (count > maximum) {70                    maximum = count;71                    maximum_index = bin_index;72                }73            }74            return maximum_index;75        }7677        pub fn write(78            self: *const Self,79            writer: *std.Io.Writer,80            caption: []const u8,81            skip_zero: bool,82        ) std.Io.Writer.Error!void {83            try writer.print("\n{s} [{d}, modal idx {d}]\n", .{84                caption,85                capacity,86                self.modalBinIndex(),87            });88            const first = self.firstNonzero();89            if (first == capacity) return;90            const last = self.lastNonzero();91            for (first..last + 1) |bin_index| {92                const count = self.counts[bin_index];93                if (!skip_zero or count != 0) {94                    try writer.print(" {d:3}: {d}\n", .{ bin_index, count });95                }96            }97        }9899        pub fn reset(self: *Self) void {100            @memset(&self.counts, 0);101        }102103        fn index(bin_value: anytype) usize {104            const T = @TypeOf(bin_value);105            comptime switch (@typeInfo(T)) {106                .int, .comptime_int => {},107                else => @compileError("bin index must be an integer"),108            };109            switch (@typeInfo(T)) {110                .int => |info| if (info.signedness == .signed) {111                    std.debug.assert(bin_value >= 0);112                },113                .comptime_int => std.debug.assert(bin_value >= 0),114                else => unreachable,115            }116            const bin_index: usize = @intCast(bin_value);117            std.debug.assert(bin_index < capacity);118            return bin_index;119        }120121        fn countValue(count: u32) Count {122            if (@bitSizeOf(Count) >= @bitSizeOf(u32)) return @intCast(count);123            return @truncate(count);124        }125    };126}127128pub const Stats = struct {129    count_value: i64 = 0,130    minimum: f32 = std.math.floatMax(f32),131    maximum: f32 = -std.math.floatMax(f32),132    sum_log: f64 = 0,133    moment_1: f64 = 0,134    moment_2: f64 = 0,135    moment_3: f64 = 0,136    moment_4: f64 = 0,137138    pub const max_count: i64 = 3_037_000_499;139    pub const no_count: u8 = 1;140    pub const no_mean_sd: u8 = 2;141    pub const no_min_max: u8 = 4;142    pub const no_skew_kurt: u8 = 8;143    pub const no_geometric_mean: u8 = 16;144    pub const all_exclusions: u8 = no_count |145        no_mean_sd |146        no_min_max |147        no_skew_kurt |148        no_geometric_mean;149150    pub fn notify(self: *@This(), value: f32) void {151        std.debug.assert(self.count_value >= 0);152        std.debug.assert(self.count_value < max_count);153        self.count_value += 1;154        self.minimum = @min(self.minimum, value);155        self.maximum = @max(self.maximum, value);156        self.sum_log += @log(@as(f64, value));157158        const count_f64: f64 = @floatFromInt(self.count_value);159        const delta = @as(f64, value) - self.moment_1;160        const delta_div_count = delta / count_f64;161        const delta2_count_minus_1_div_count = delta * (count_f64 - 1) * delta_div_count;162        const count_polynomial = count_f64 * count_f64 - 3 * count_f64 + 3;163        self.moment_1 += delta_div_count;164        self.moment_4 += delta_div_count *165            (delta_div_count *166                (delta2_count_minus_1_div_count * count_polynomial + 6 * self.moment_2) -167                4 * self.moment_3);168        self.moment_3 += delta_div_count *169            (delta2_count_minus_1_div_count * (count_f64 - 2) - 3 * self.moment_2);170        self.moment_2 += delta2_count_minus_1_div_count;171    }172173    pub fn assimilate(self: *@This(), other: *const @This()) void {174        std.debug.assert(self.count_value >= 0);175        std.debug.assert(other.count_value >= 0);176        std.debug.assert(self.count_value <= max_count - other.count_value);177        const total_count = self.count_value + other.count_value;178        if (total_count == 0) return;179180        self.minimum = @min(self.minimum, other.minimum);181        self.maximum = @max(self.maximum, other.maximum);182        self.sum_log += other.sum_log;183184        const own_count: f64 = @floatFromInt(self.count_value);185        const other_count: f64 = @floatFromInt(other.count_value);186        const total_count_f64: f64 = @floatFromInt(total_count);187        const product = own_count * other_count;188        const count_squared = own_count * own_count;189        const other_count_squared = other_count * other_count;190        const total_count_squared = total_count_f64 * total_count_f64;191        const total_count_cubed = total_count_squared * total_count_f64;192        const inverse_total_count = 1 / total_count_f64;193        const inverse_total_count_squared = 1 / total_count_squared;194195        const delta = other.moment_1 - self.moment_1;196        const delta_squared = delta * delta;197        const delta_cubed = delta * delta_squared;198        const delta_fourth = delta_squared * delta_squared;199200        self.moment_1 = (own_count * self.moment_1 + other_count * other.moment_1) * inverse_total_count;201        const new_moment_2 = self.moment_2 + other.moment_2 +202            delta_squared * product * inverse_total_count;203        const new_moment_3 = self.moment_3 + other.moment_3 +204            delta_cubed * product * (own_count - other_count) * inverse_total_count_squared +205            3 * delta * (own_count * other.moment_2 - other_count * self.moment_2) * inverse_total_count;206        self.moment_4 += other.moment_4 +207            delta_fourth * product * (count_squared - product + other_count_squared) /208                total_count_cubed +209            6 * delta_squared *210                (count_squared * other.moment_2 + other_count_squared * self.moment_2) *211                inverse_total_count_squared +212            4 * delta * (own_count * other.moment_3 - other_count * self.moment_3) *213                inverse_total_count;214        self.moment_2 = new_moment_2;215        self.moment_3 = new_moment_3;216        self.count_value = total_count;217    }218219    pub fn count(self: *const @This()) i64 {220        return self.count_value;221    }222223    pub fn min(self: *const @This()) f32 {224        return self.minimum;225    }226227    pub fn max(self: *const @This()) f32 {228        return self.maximum;229    }230231    pub fn geometricMean(self: *const @This()) f64 {232        if (self.count_value == 0) return 0;233        return @exp(self.sum_log / @as(f64, @floatFromInt(self.count_value)));234    }235236    pub fn mean(self: *const @This()) f64 {237        return self.moment_1;238    }239240    pub fn sampleVariance(self: *const @This()) f64 {241        if (self.count_value == 0) return 0;242        return self.moment_2 / @as(f64, @floatFromInt(self.count_value));243    }244245    pub fn variance(self: *const @This()) f64 {246        if (self.count_value == 0) return 0;247        if (self.count_value == 1) return self.moment_2;248        return self.moment_2 / @as(f64, @floatFromInt(self.count_value - 1));249    }250251    pub fn standardDeviation(self: *const @This()) f64 {252        return @sqrt(self.variance());253    }254255    pub fn sampleSkewness(self: *const @This()) f64 {256        if (@abs(self.moment_2) < 1e-7) return 0;257        const count_f64: f64 = @floatFromInt(self.count_value);258        return self.moment_3 * @sqrt(count_f64) /259            std.math.pow(f64, self.moment_2, 1.5);260    }261262    pub fn skewness(self: *const @This()) f64 {263        if (self.count_value == 0) return 0;264        const count_f64: f64 = @floatFromInt(self.count_value);265        const ratio = (count_f64 - 1) / count_f64;266        return self.sampleSkewness() * std.math.pow(f64, ratio, 1.5);267    }268269    pub fn sampleKurtosis(self: *const @This()) f64 {270        if (@abs(self.moment_2) < 1e-7) return 0;271        const count_f64: f64 = @floatFromInt(self.count_value);272        return self.moment_4 * count_f64 / (self.moment_2 * self.moment_2);273    }274275    pub fn kurtosis(self: *const @This()) f64 {276        if (self.count_value == 0) return 0;277        const count_f64: f64 = @floatFromInt(self.count_value);278        const ratio = (count_f64 - 1) / count_f64;279        return self.sampleKurtosis() * ratio * ratio;280    }281282    pub fn mu1(self: *const @This()) f64 {283        std.debug.assert(self.count_value != 0);284        return self.moment_1;285    }286287    pub fn mu2(self: *const @This()) f64 {288        return self.centralMoment(self.moment_2);289    }290291    pub fn mu3(self: *const @This()) f64 {292        return self.centralMoment(self.moment_3);293    }294295    pub fn mu4(self: *const @This()) f64 {296        return self.centralMoment(self.moment_4);297    }298299    pub fn write(300        self: *const @This(),301        writer: *std.Io.Writer,302        exclude: u8,303    ) std.Io.Writer.Error!void {304        if (self.count_value == 0) return writer.writeAll("(none)");305        if (exclude & no_count == 0) {306            try writer.print("Count={d:9} ", .{@as(u64, @intCast(self.count_value))});307        }308        if (exclude & no_mean_sd == 0) {309            try writer.writeAll("Mean=");310            try writeCScientific(writer, self.mean(), 10, 3);311            try writer.writeAll(" SD=");312            try writeCScientific(writer, self.standardDeviation(), 8, 2);313            try writer.writeByte(' ');314        }315        if (exclude & no_min_max == 0) {316            try writer.writeAll("Min=");317            try writeCScientific(writer, self.minimum, 10, 3);318            try writer.writeAll(" Max=");319            try writeCScientific(writer, self.maximum, 10, 3);320            try writer.writeByte(' ');321        }322        if (exclude & no_skew_kurt == 0) {323            try writer.print("Skew={d:5.2} Kurt={d:7.2} ", .{324                self.skewness(),325                self.kurtosis(),326            });327        }328        if (exclude & no_geometric_mean == 0) {329            try writer.print("GeoMean={d:9.6} ", .{self.geometricMean()});330        }331    }332333    pub fn reset(self: *@This()) void {334        self.* = .{};335    }336337    fn centralMoment(self: *const @This(), moment: f64) f64 {338        std.debug.assert(self.count_value != 0);339        return moment / @as(f64, @floatFromInt(self.count_value));340    }341};342343fn writeCScientific(344    writer: *std.Io.Writer,345    value: anytype,346    width: usize,347    precision: usize,348) std.Io.Writer.Error!void {349    var rendered_buffer: [std.fmt.float.min_buffer_size]u8 = undefined;350    const rendered = std.fmt.float.render(&rendered_buffer, value, .{351        .mode = .scientific,352        .precision = precision,353    }) catch unreachable;354    const exponent_marker = std.mem.lastIndexOfScalar(u8, rendered, 'e') orelse {355        return writer.alignBuffer(rendered, width, .right, ' ');356    };357    const exponent = std.fmt.parseInt(i32, rendered[exponent_marker + 1 ..], 10) catch unreachable;358    var exponent_buffer: [16]u8 = undefined;359    const exponent_digits = std.fmt.bufPrint(&exponent_buffer, "{d}", .{@abs(exponent)}) catch unreachable;360    const zero_count = if (exponent_digits.len < 2) 2 - exponent_digits.len else 0;361    const output_length = exponent_marker + 2 + zero_count + exponent_digits.len;362    if (output_length < width) try writer.splatByteAll(' ', width - output_length);363    try writer.writeAll(rendered[0..exponent_marker]);364    try writer.writeByte('e');365    try writer.writeByte(if (exponent < 0) '-' else '+');366    try writer.splatByteAll('0', zero_count);367    try writer.writeAll(exponent_digits);368}369370fn requireBins(comptime capacity: usize, comptime Count: type) void {371    if (capacity == 0) @compileError("bins require a nonzero capacity");372    switch (@typeInfo(Count)) {373        .int => |info| if (info.signedness != .unsigned) {374            @compileError("bin counts must use an unsigned integer type");375        },376        else => @compileError("bin counts must use an unsigned integer type"),377    }378}379380test "fixed bins notify, assimilate, summarize, and reset" {381    const Histogram = Bins(6);382    var bins = Histogram{};383    try std.testing.expectEqual(@as(usize, 6), bins.firstNonzero());384    try std.testing.expectEqual(@as(usize, 0), bins.lastNonzero());385    try std.testing.expectEqual(@as(usize, 0), bins.numNonzero());386    try std.testing.expectEqual(@as(usize, 0), bins.modalBinIndex());387388    bins.notify(@as(i8, 4));389    bins.incrementBy(@as(u16, 2), 3);390    bins.incrementBy(4, 2);391    var other = Histogram{};392    other.incrementBy(3, 2);393    other.notify(4);394    bins.assimilate(&other);395    try std.testing.expectEqual(@as(u32, 3), bins.bin(2));396    try std.testing.expectEqual(@as(u32, 2), bins.bin(3));397    try std.testing.expectEqual(@as(u32, 4), bins.bin(4));398    try std.testing.expectEqual(@as(usize, 2), bins.firstNonzero());399    try std.testing.expectEqual(@as(usize, 4), bins.lastNonzero());400    try std.testing.expectEqual(@as(usize, 3), bins.numNonzero());401    try std.testing.expectEqual(@as(usize, 4), bins.modalBinIndex());402403    bins.resetBin(4);404    try std.testing.expectEqual(@as(usize, 2), bins.modalBinIndex());405    bins.reset();406    try std.testing.expectEqual(@as(usize, 0), bins.numNonzero());407}408409test "typed fixed bins preserve unsigned count wrapping" {410    var bins = TypedBins(2, u8){};411    bins.incrementBy(1, 300);412    try std.testing.expectEqual(@as(u8, 44), bins.bin(1));413    bins.incrementBy(1, 212);414    try std.testing.expectEqual(@as(u8, 0), bins.bin(1));415}416417test "fixed bins writer preserves Highway text layout" {418    var bins = Bins(6){};419    bins.incrementBy(1, 2);420    bins.incrementBy(3, 5);421    var buffer: [256]u8 = undefined;422    var writer = std.Io.Writer.fixed(&buffer);423    try bins.write(&writer, "latency", false);424    try std.testing.expectEqualStrings(425        "\nlatency [6, modal idx 3]\n   1: 2\n   2: 0\n   3: 5\n",426        writer.buffered(),427    );428429    writer = std.Io.Writer.fixed(&buffer);430    try bins.write(&writer, "latency", true);431    try std.testing.expectEqualStrings(432        "\nlatency [6, modal idx 3]\n   1: 2\n   3: 5\n",433        writer.buffered(),434    );435}436437test "online statistics handle empty, singleton, and reset states" {438    var stats = Stats{};439    try std.testing.expectEqual(@as(i64, 0), stats.count());440    try std.testing.expectEqual(std.math.floatMax(f32), stats.min());441    try std.testing.expectEqual(-std.math.floatMax(f32), stats.max());442    try std.testing.expectEqual(@as(f64, 0), stats.geometricMean());443    try std.testing.expectEqual(@as(f64, 0), stats.mean());444    try std.testing.expectEqual(@as(f64, 0), stats.variance());445    try std.testing.expectEqual(@as(f64, 0), stats.standardDeviation());446    try std.testing.expectEqual(@as(f64, 0), stats.skewness());447    try std.testing.expectEqual(@as(f64, 0), stats.kurtosis());448    var buffer: [32]u8 = undefined;449    var writer = std.Io.Writer.fixed(&buffer);450    try stats.write(&writer, 0);451    try std.testing.expectEqualStrings("(none)", writer.buffered());452453    stats.notify(4);454    try std.testing.expectEqual(@as(i64, 1), stats.count());455    try std.testing.expectEqual(@as(f32, 4), stats.min());456    try std.testing.expectEqual(@as(f32, 4), stats.max());457    try std.testing.expectEqual(@as(f64, 4), stats.mean());458    try std.testing.expectEqual(@as(f64, 0), stats.variance());459    try std.testing.expectEqual(@as(f64, 0), stats.skewness());460    try std.testing.expectEqual(@as(f64, 0), stats.kurtosis());461    try std.testing.expectApproxEqRel(@as(f64, 4), stats.geometricMean(), 1e-15);462    stats.reset();463    try std.testing.expectEqual(@as(i64, 0), stats.count());464}465466test "online statistics assimilate partitions" {467    const values = [_]f32{ 1, 2, 2, 3, 5, 8, 13, 21 };468    var sequential = Stats{};469    for (values) |value| sequential.notify(value);470    var left = Stats{};471    for (values[0..3]) |value| left.notify(value);472    var right = Stats{};473    for (values[3..]) |value| right.notify(value);474    left.assimilate(&right);475476    try std.testing.expectEqual(sequential.count(), left.count());477    try std.testing.expectEqual(sequential.min(), left.min());478    try std.testing.expectEqual(sequential.max(), left.max());479    try std.testing.expectApproxEqRel(sequential.geometricMean(), left.geometricMean(), 1e-15);480    try std.testing.expectApproxEqRel(sequential.mean(), left.mean(), 1e-15);481    try std.testing.expectApproxEqRel(sequential.variance(), left.variance(), 1e-14);482    try std.testing.expectApproxEqRel(sequential.skewness(), left.skewness(), 1e-14);483    try std.testing.expectApproxEqRel(sequential.kurtosis(), left.kurtosis(), 1e-14);484    try std.testing.expectApproxEqRel(sequential.mu1(), left.mu1(), 1e-15);485    try std.testing.expectApproxEqRel(sequential.mu2(), left.mu2(), 1e-14);486    try std.testing.expectApproxEqRel(sequential.mu3(), left.mu3(), 1e-14);487    try std.testing.expectApproxEqRel(sequential.mu4(), left.mu4(), 1e-14);488}489490test "online statistics match the pinned Highway oracle" {491    const values = [_]f32{ 1, 2, 2, 3, 5, 8, 13, 21 };492    var stats = Stats{};493    for (values) |value| stats.notify(value);494495    try std.testing.expectEqual(@as(i64, 8), stats.count());496    try std.testing.expectEqual(@as(f32, 1), stats.min());497    try std.testing.expectEqual(@as(f32, 21), stats.max());498    try std.testing.expectEqual(@as(f64, 0x1.17295561888e1p+2), stats.geometricMean());499    try std.testing.expectEqual(@as(f64, 0x1.b8p+2), stats.mean());500    try std.testing.expectEqual(@as(f64, 0x1.52ep+5), stats.sampleVariance());501    try std.testing.expectEqual(@as(f64, 0x1.8349249249249p+5), stats.variance());502    try std.testing.expectEqual(@as(f64, 0x1.bd4c4807664aep+2), stats.standardDeviation());503    try std.testing.expectApproxEqAbs(@as(f64, 0x1.27fe45be6a283p+0), stats.sampleSkewness(), 4e-15);504    try std.testing.expectApproxEqAbs(@as(f64, 0x1.e488990369159p-1), stats.skewness(), 4e-15);505    try std.testing.expectApproxEqAbs(@as(f64, 0x1.86564df4935ecp+1), stats.sampleKurtosis(), 4e-15);506    try std.testing.expectApproxEqAbs(@as(f64, 0x1.2ada13af40d48p+1), stats.kurtosis(), 4e-15);507    try std.testing.expectEqual(@as(f64, 0x1.b8p+2), stats.mu1());508    try std.testing.expectEqual(@as(f64, 0x1.52ep+5), stats.mu2());509    try std.testing.expectEqual(@as(f64, 0x1.3ec3p+8), stats.mu3());510    try std.testing.expectEqual(@as(f64, 0x1.55fc8dp+12), stats.mu4());511512    var buffer: [300]u8 = undefined;513    var writer = std.Io.Writer.fixed(&buffer);514    try stats.write(&writer, 0);515    try std.testing.expectEqualStrings(516        "Count=        8 Mean= 6.875e+00 SD=6.96e+00 Min= 1.000e+00 Max= 2.100e+01 Skew= 0.95 Kurt=   2.33 GeoMean= 4.361898 ",517        writer.buffered(),518    );519520    writer = std.Io.Writer.fixed(&buffer);521    try stats.write(&writer, Stats.no_count | Stats.no_min_max | Stats.no_geometric_mean);522    try std.testing.expectEqualStrings(523        "Mean= 6.875e+00 SD=6.96e+00 Skew= 0.95 Kurt=   2.33 ",524        writer.buffered(),525    );526527    writer = std.Io.Writer.fixed(&buffer);528    try stats.write(&writer, Stats.all_exclusions | 0x80);529    try std.testing.expectEqual(@as(usize, 0), writer.buffered().len);530}531532test "online statistics assimilate empty states" {533    var empty = Stats{};534    var full = Stats{};535    for ([_]f32{ 2, 3, 7 }) |value| full.notify(value);536    const expected = full;537    full.assimilate(&empty);538    try std.testing.expectEqualDeep(expected, full);539    empty.assimilate(&expected);540    try std.testing.expectEqualDeep(expected, empty);541}542543test "online statistics support self assimilation" {544    const values = [_]f32{ 1, 3, 7, 9 };545    var doubled = Stats{};546    for (values) |value| doubled.notify(value);547    doubled.assimilate(&doubled);548    var sequential = Stats{};549    for (0..2) |_| {550        for (values) |value| sequential.notify(value);551    }552    try std.testing.expectEqual(sequential.count(), doubled.count());553    try std.testing.expectEqual(sequential.min(), doubled.min());554    try std.testing.expectEqual(sequential.max(), doubled.max());555    try std.testing.expectApproxEqRel(sequential.geometricMean(), doubled.geometricMean(), 1e-15);556    try std.testing.expectApproxEqRel(sequential.mean(), doubled.mean(), 1e-15);557    try std.testing.expectApproxEqRel(sequential.mu2(), doubled.mu2(), 1e-15);558    try std.testing.expectApproxEqRel(sequential.mu3(), doubled.mu3(), 1e-15);559    try std.testing.expectApproxEqRel(sequential.mu4(), doubled.mu4(), 1e-15);560}561562test "geometric mean retains logarithm domain behavior" {563    var zero = Stats{};564    zero.notify(0);565    zero.notify(4);566    try std.testing.expectEqual(@as(f64, 0), zero.geometricMean());567    var negative = Stats{};568    negative.notify(-1);569    try std.testing.expect(std.math.isNan(negative.geometricMean()));570}