1// Ported from musl, which is licensed under the MIT license:
  2// https://git.musl-libc.org/cgit/musl/tree/COPYRIGHT
  3//
  4// https://git.musl-libc.org/cgit/musl/tree/src/complex/csqrtf.c
  5// https://git.musl-libc.org/cgit/musl/tree/src/complex/csqrt.c
  6
  7const std = @import("../../std.zig");
  8const testing = std.testing;
  9const math = std.math;
 10const cmath = math.complex;
 11const Complex = cmath.Complex;
 12
 13/// Returns the square root of z. The real and imaginary parts of the result have the same sign
 14/// as the imaginary part of z.
 15pub fn sqrt(z: anytype) Complex(@TypeOf(z.re, z.im)) {
 16    const T = @TypeOf(z.re, z.im);
 17
 18    return switch (T) {
 19        f32 => sqrt32(z),
 20        f64 => sqrt64(z),
 21        else => @compileError("sqrt not implemented for " ++ @typeName(T)),
 22    };
 23}
 24
 25fn sqrt32(z: Complex(f32)) Complex(f32) {
 26    const x = z.re;
 27    const y = z.im;
 28
 29    if (x == 0 and y == 0) {
 30        return Complex(f32).init(0, y);
 31    }
 32    if (math.isInf(y)) {
 33        return Complex(f32).init(math.inf(f32), y);
 34    }
 35    if (math.isNan(x)) {
 36        // raise invalid if y is not nan
 37        const t = (y - y) / (y - y);
 38        return Complex(f32).init(x, t);
 39    }
 40    if (math.isInf(x)) {
 41        // sqrt(inf + i nan)    = inf + nan i
 42        // sqrt(inf + iy)       = inf + i0
 43        // sqrt(-inf + i nan)   = nan +- inf i
 44        // sqrt(-inf + iy)      = 0 + inf i
 45        if (math.signbit(x)) {
 46            return Complex(f32).init(@abs(y - y), math.copysign(x, y));
 47        } else {
 48            return Complex(f32).init(x, math.copysign(y - y, y));
 49        }
 50    }
 51
 52    // y = nan special case is handled fine below
 53
 54    // double-precision avoids overflow with correct rounding.
 55    const dx = @as(f64, x);
 56    const dy = @as(f64, y);
 57
 58    if (dx >= 0) {
 59        const t = @sqrt((dx + math.hypot(dx, dy)) * 0.5);
 60        return Complex(f32).init(
 61            @as(f32, @floatCast(t)),
 62            @as(f32, @floatCast(dy / (2.0 * t))),
 63        );
 64    } else {
 65        const t = @sqrt((-dx + math.hypot(dx, dy)) * 0.5);
 66        return Complex(f32).init(
 67            @as(f32, @floatCast(@abs(y) / (2.0 * t))),
 68            @as(f32, @floatCast(math.copysign(t, y))),
 69        );
 70    }
 71}
 72
 73fn sqrt64(z: Complex(f64)) Complex(f64) {
 74    // may encounter overflow for im,re >= DBL_MAX / (1 + sqrt(2))
 75    const threshold = 0x1.a827999fcef32p+1022;
 76
 77    var x = z.re;
 78    var y = z.im;
 79
 80    if (x == 0 and y == 0) {
 81        return Complex(f64).init(0, y);
 82    }
 83    if (math.isInf(y)) {
 84        return Complex(f64).init(math.inf(f64), y);
 85    }
 86    if (math.isNan(x)) {
 87        // raise invalid if y is not nan
 88        const t = (y - y) / (y - y);
 89        return Complex(f64).init(x, t);
 90    }
 91    if (math.isInf(x)) {
 92        // sqrt(inf + i nan)    = inf + nan i
 93        // sqrt(inf + iy)       = inf + i0
 94        // sqrt(-inf + i nan)   = nan +- inf i
 95        // sqrt(-inf + iy)      = 0 + inf i
 96        if (math.signbit(x)) {
 97            return Complex(f64).init(@abs(y - y), math.copysign(x, y));
 98        } else {
 99            return Complex(f64).init(x, math.copysign(y - y, y));
100        }
101    }
102
103    // y = nan special case is handled fine below
104
105    // scale to avoid overflow
106    var scale = false;
107    if (@abs(x) >= threshold or @abs(y) >= threshold) {
108        x *= 0.25;
109        y *= 0.25;
110        scale = true;
111    }
112
113    var result: Complex(f64) = undefined;
114    if (x >= 0) {
115        const t = @sqrt((x + math.hypot(x, y)) * 0.5);
116        result = Complex(f64).init(t, y / (2.0 * t));
117    } else {
118        const t = @sqrt((-x + math.hypot(x, y)) * 0.5);
119        result = Complex(f64).init(@abs(y) / (2.0 * t), math.copysign(t, y));
120    }
121
122    if (scale) {
123        result.re *= 2;
124        result.im *= 2;
125    }
126
127    return result;
128}
129
130test sqrt32 {
131    const epsilon = math.floatEps(f32);
132    const a = Complex(f32).init(5, 3);
133    const c = sqrt(a);
134
135    try testing.expectApproxEqAbs(2.3271174, c.re, epsilon);
136    try testing.expectApproxEqAbs(0.6445742, c.im, epsilon);
137}
138
139test sqrt64 {
140    const epsilon = math.floatEps(f64);
141    const a = Complex(f64).init(5, 3);
142    const c = sqrt(a);
143
144    try testing.expectApproxEqAbs(2.3271175190399496, c.re, epsilon);
145    try testing.expectApproxEqAbs(0.6445742373246469, c.im, epsilon);
146}