diff --git a/src/numz.zig b/src/numz.zig index 0f02e79..bf40bd9 100644 --- a/src/numz.zig +++ b/src/numz.zig @@ -253,16 +253,39 @@ pub fn Matrix(rows: comptime_int, cols: comptime_int, T: type) type { inline for (0..result.cols) |j| { var sum: T = 0; inline for (0..self.cols) |k| { - if (@typeInfo(T) == .float) - sum = @mulAdd(T, self.mat[i][k], other.mat[k][j], sum) - else - sum += self.mat[i][k] * other.mat[k][j]; + switch (@typeInfo(T)) { + .float => sum = @mulAdd(T, self.mat[i][k], other.mat[k][j], sum), + else => sum += self.mat[i][k] * other.mat[k][j], + } } result.mat[i][j] = sum; } } return result; } + + pub fn normSquared(self: Self) T { + var n: T = 0; + inline for (self.mat) |row| { + inline for (row) |e| { + switch (@typeInfo(T)) { + .float => n = @mulAdd(T, e, e, n), + else => n += e * e, + } + } + } + return n; + } + + pub const norm = + switch (@typeInfo(T)) { + .float => struct { + fn norm(self: Self) T { + return @sqrt(self.normSquared()); + } + }.norm, + else => @compileError("Connot take a square root on an integer matrix"), + }; }; } @@ -467,3 +490,11 @@ test "stack scalar cols" { const expected: [1][6]f32 = .{.{ 1, 2, 3, 4, 5, 6 }}; try std.testing.expectEqual(expected, m2.mat); } + +test "norm" { + const int = stackRows(u32, .{ 3, 4 }); + try std.testing.expectEqual(25, int.normSquared()); + + const float = stackRows(f64, .{ 3, 4 }); + try std.testing.expectEqual(5, float.norm()); +}