From 87863a834b074311d70ec4dd3b10c74705fff54d Mon Sep 17 00:00:00 2001
From: Chris Boesch <48591413+chrboesch@users.noreply.github.com>
Date: Tue, 26 Nov 2024 14:03:48 +0100
Subject: [PATCH] std.math.complex: Add squared magnitude function (#21998)

---
 lib/std/math/complex.zig | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/lib/std/math/complex.zig b/lib/std/math/complex.zig
index 8398d0a69a..639474fb2a 100644
--- a/lib/std/math/complex.zig
+++ b/lib/std/math/complex.zig
@@ -115,6 +115,10 @@ pub fn Complex(comptime T: type) type {
         pub fn magnitude(self: Self) T {
             return @sqrt(self.re * self.re + self.im * self.im);
         }
+
+        pub fn squaredMagnitude(self: Self) T {
+            return self.re * self.re + self.im * self.im;
+        }
     };
 }
 
@@ -189,6 +193,13 @@ test "magnitude" {
     try testing.expect(math.approxEqAbs(f32, c, 5.83095, epsilon));
 }
 
+test "squaredMagnitude" {
+    const a = Complex(f32).init(5, 3);
+    const c = a.squaredMagnitude();
+
+    try testing.expect(math.approxEqAbs(f32, c, math.pow(f32, a.magnitude(), 2), epsilon));
+}
+
 test {
     _ = @import("complex/abs.zig");
     _ = @import("complex/acosh.zig");