139 lines
3.8 KiB
Zig
139 lines
3.8 KiB
Zig
const std = @import("../../index.zig");
|
|
const debug = std.debug;
|
|
const math = std.math;
|
|
const cmath = math.complex;
|
|
const Complex = cmath.Complex;
|
|
|
|
pub fn sqrt(z: var) @typeOf(z) {
|
|
const T = @typeOf(z.re);
|
|
|
|
return switch (T) {
|
|
f32 => sqrt32(z),
|
|
f64 => sqrt64(z),
|
|
else => @compileError("sqrt not implemented for " ++ @typeName(T)),
|
|
};
|
|
}
|
|
|
|
fn sqrt32(z: Complex(f32)) Complex(f32) {
|
|
const x = z.re;
|
|
const y = z.im;
|
|
|
|
if (x == 0 and y == 0) {
|
|
return Complex(f32).new(0, y);
|
|
}
|
|
if (math.isInf(y)) {
|
|
return Complex(f32).new(math.inf(f32), y);
|
|
}
|
|
if (math.isNan(x)) {
|
|
// raise invalid if y is not nan
|
|
const t = (y - y) / (y - y);
|
|
return Complex(f32).new(x, t);
|
|
}
|
|
if (math.isInf(x)) {
|
|
// sqrt(inf + i nan) = inf + nan i
|
|
// sqrt(inf + iy) = inf + i0
|
|
// sqrt(-inf + i nan) = nan +- inf i
|
|
// sqrt(-inf + iy) = 0 + inf i
|
|
if (math.signbit(x)) {
|
|
return Complex(f32).new(math.fabs(x - y), math.copysign(f32, x, y));
|
|
} else {
|
|
return Complex(f32).new(x, math.copysign(f32, y - y, y));
|
|
}
|
|
}
|
|
|
|
// y = nan special case is handled fine below
|
|
|
|
// double-precision avoids overflow with correct rounding.
|
|
const dx = f64(x);
|
|
const dy = f64(y);
|
|
|
|
if (dx >= 0) {
|
|
const t = math.sqrt((dx + math.hypot(f64, dx, dy)) * 0.5);
|
|
return Complex(f32).new(
|
|
@floatCast(f32, t),
|
|
@floatCast(f32, dy / (2.0 * t)),
|
|
);
|
|
} else {
|
|
const t = math.sqrt((-dx + math.hypot(f64, dx, dy)) * 0.5);
|
|
return Complex(f32).new(
|
|
@floatCast(f32, math.fabs(y) / (2.0 * t)),
|
|
@floatCast(f32, math.copysign(f64, t, y)),
|
|
);
|
|
}
|
|
}
|
|
|
|
fn sqrt64(z: Complex(f64)) Complex(f64) {
|
|
// may encounter overflow for im,re >= DBL_MAX / (1 + sqrt(2))
|
|
const threshold = 0x1.a827999fcef32p+1022;
|
|
|
|
var x = z.re;
|
|
var y = z.im;
|
|
|
|
if (x == 0 and y == 0) {
|
|
return Complex(f64).new(0, y);
|
|
}
|
|
if (math.isInf(y)) {
|
|
return Complex(f64).new(math.inf(f64), y);
|
|
}
|
|
if (math.isNan(x)) {
|
|
// raise invalid if y is not nan
|
|
const t = (y - y) / (y - y);
|
|
return Complex(f64).new(x, t);
|
|
}
|
|
if (math.isInf(x)) {
|
|
// sqrt(inf + i nan) = inf + nan i
|
|
// sqrt(inf + iy) = inf + i0
|
|
// sqrt(-inf + i nan) = nan +- inf i
|
|
// sqrt(-inf + iy) = 0 + inf i
|
|
if (math.signbit(x)) {
|
|
return Complex(f64).new(math.fabs(x - y), math.copysign(f64, x, y));
|
|
} else {
|
|
return Complex(f64).new(x, math.copysign(f64, y - y, y));
|
|
}
|
|
}
|
|
|
|
// y = nan special case is handled fine below
|
|
|
|
// scale to avoid overflow
|
|
var scale = false;
|
|
if (math.fabs(x) >= threshold or math.fabs(y) >= threshold) {
|
|
x *= 0.25;
|
|
y *= 0.25;
|
|
scale = true;
|
|
}
|
|
|
|
var result: Complex(f64) = undefined;
|
|
if (x >= 0) {
|
|
const t = math.sqrt((x + math.hypot(f64, x, y)) * 0.5);
|
|
result = Complex(f64).new(t, y / (2.0 * t));
|
|
} else {
|
|
const t = math.sqrt((-x + math.hypot(f64, x, y)) * 0.5);
|
|
result = Complex(f64).new(math.fabs(y) / (2.0 * t), math.copysign(f64, t, y));
|
|
}
|
|
|
|
if (scale) {
|
|
result.re *= 2;
|
|
result.im *= 2;
|
|
}
|
|
|
|
return result;
|
|
}
|
|
|
|
const epsilon = 0.0001;
|
|
|
|
test "complex.csqrt32" {
|
|
const a = Complex(f32).new(5, 3);
|
|
const c = sqrt(a);
|
|
|
|
debug.assert(math.approxEq(f32, c.re, 2.327117, epsilon));
|
|
debug.assert(math.approxEq(f32, c.im, 0.644574, epsilon));
|
|
}
|
|
|
|
test "complex.csqrt64" {
|
|
const a = Complex(f64).new(5, 3);
|
|
const c = sqrt(a);
|
|
|
|
debug.assert(math.approxEq(f64, c.re, 2.3271175190399496, epsilon));
|
|
debug.assert(math.approxEq(f64, c.im, 0.6445742373246469, epsilon));
|
|
}
|