Complex Square Roots in math.js

Math.js is an extensive mathematics library for Javascript, with support for complex numbers, matrices, and unit conversions, along with utility functions for, for example, generating random values from various distributions. There is also beauty and subtlety to correctly implementing mathematical code, something I’ve come to appreciate more and more over the course of this project. So when math.js popped up on Hacker News, I had to try passing it through Herbie.

I should add that math.js is a very well-written library. There are extensive tests (including tests for edge cases, like arc-cosine of 1); there is detailed documentation; a well-maintained website; a build-bot; and a responsive author.

Math.js contains many functions, and I’m far from passing them all through Herbie. For now, Herbie’s input must be a single expression, written in a stilted Scheme dialect. I manually translated from math.js’s arithmetic functions into this format; the functions which include complex number multiplication, exponentiation, and roots. Each math.js function resulted in multiple Herbie expressions: functions of a complex number separately return a real and an imaginary part, two separate Herbie expressions. All told, I extracted 28 floating-point expressions.

After the tedious work of translating Javascript to Scheme, Herbie runs with little manual intervention. Herbie is designed to require minimal expert knowledge, so there aren’t knobs to tweak or prompts to respond to.

At first, Herbie wasn’t finding any expressions that it could improve (though it did find several to be inaccurate). Discouraged, I went on to other things, but a few weeks later Alex, who’s been doing some fantastic heavy lifting in the Herbie internals, made a big improvement to our simplification engine. Simplification rearranges expressions to cancel like terms (turning a sums like \(a + b - a + c - b\) into the simpler \(c\)); the old version couldn’t cancel pairs like \(-x · x\) and \(x^2\). With this fix, Herbie found an improved way of calculating complex square roots.

Given a complex number \(x + i y\), its square root has real part \[ \frac1{\sqrt2} \sqrt{\sqrt{x^2 + y^2} + x}, \] and imaginary part \[ \frac1{\sqrt2} \sqrt{\sqrt{x^2 + y^2} - x}, \] with the signs for the outer square root dependent on the signs of \(x\) and \(y\). In math.js, this formula is implemented by the following Javascript. Note that \(x\) and \(y\) become x.re and x.im, because complex numbers are objects in math.js² [² Why not just start with the code, instead of jumping between math and code? Mathematics is a much more compact notation, so it is easier to read and understand; the code contains several branches, and the control flow is hard to follow (but they are unimportant for the main ideas); and, comparing the real-number and floating-point semantics of an expression is central to this post, so it is best not to pick a notation which already assumes floating-point semantics and computability.].

var r = Math.sqrt(x.re * x.re + x.im * x.im);
if (x.im >= 0) {
  return new Complex(
    0.5 * Math.sqrt(2.0 * (r + x.re)),
    0.5 * Math.sqrt(2.0 * (r - x.re))
  );
}
else {
  return new Complex(
    0.5 * Math.sqrt(2.0 * (r + x.re)),
    -0.5 * Math.sqrt(2.0 * (r - x.re))
  );
}

This is a mathematician’s definition, and works when \(x\) and \(y\) are infinite-precision real numbers. But when this expression is evaluated with floating-point operators and values, it can yield wildly inaccurate results.

As I discussed in Floating Point Guarantees, every time a basic floating point operation (like addition, multiplication, or square root) is executed, the ideal real number result has to be rounded back to a value representable as a float. So, the value of \(x² + y²\) is rounded, and if \(x\) is large and \(y\) is small, \(y\) is ignored or truncated. Then, a square root is taken (leaving approximately \(\sqrt{x^2 + \epsilon} \approx |x| + \epsilon / 2 |x|\)) and \(x\) is added on. If \(x\) is negative, the sum of \(x\) and approximately \(|x|\) subtracts two almost-opposite terms, leaving behind only a very small value.

The result is small, but it is computed by subtracting two large values. When these values were computed, they were rounded, and rounding introduces error proportional in size to the rounded value. So, when large values are computed, rounding introduces a (relatively) large error; and when the result is small, this error has an outsize effect, and can lead to the result being completely inaccurate.

We can get an estimate of the effect of rounding like so. Suppose \(y\) is small, and \(x\) is large and large. Then \(\sqrt{x^2 + y^2}\) is approximately \(-x - y² / (2 x)\), and the associated rounding error is approximately \(x ε\) (for some small \(ε\); for double precision, it is approximately \(10⁻¹⁶\)). But then the \(-x\) term cancels, leaving behind \(-y² / (2 x) + x ε\) for the value inside the outer square root. The relative error is \(x ε / (y² / (2 x)) ≈ 2 (x / y)² ε\), which can be very large.

The moral of the story is that the usual expression for computing complex square roots doesn’t work as well as it could for complex numbers with negative real parts. Put in any value with large negative \(x\) and small \(y\) (for example, \(x = -10^{10}\) and \(y = 10^{-10}\)), and instead of the correct real part (which should be approximately \(5·10^{-16}\)) you get \(0\). To get the correct, more accurate answer, a different formula is necessary.

There is a standard approach to fixing problems of this sort: rearrange the expression, using the standard rules of real-number arithmetic, into a form that doesn’t trigger the same bad behavior. Herbie is designed to do this automatically.

For the complex square root, Herbie applies an error analysis to come to the same conclusion that we did: the problem is the cancellation (subtraction of near-equal numbers) of \(\sqrt{x^2 + y^2}\) and \(x\). Herbie proceeds by applying a database of rewrite rules, each basic facts about real-number arithmetic. In this case, the rewrite rule \(a + b = (a² - b²) / (a - b)\), transforming the expression into \[\frac1{\sqrt2} \sqrt{\frac{y^2}{\sqrt{x^2 + y^2} - x}}.\] While it may not be obvious, the cancellation has disappeared: when \(x\) is negative, the two values that are subtracted have opposite signs, so the subtraction doesn’t cancel anything.

Now, this new expression is very accurate for negative \(x\), but for positive \(x\) it has its own cancellation problem. The solution is to test whether \(x\) is positive or negative, and apply the appropriate expression. Herbie also does this automatically, and returns an expression with the branch already put in.

Herbie returns its output in the same Scheme dialect that it uses for input, so I manually translated its output into Javascript. I made the same modification to the imaginary part, leaving me with this code:

var r = Math.sqrt(x.re * x.re + x.im * x.im);
var re, im;

if (x.re >= 0) {
  re = 0.5 * Math.sqrt(2.0 * (r + x.re));
}
else {
  re = Math.abs(x.im) / Math.sqrt(2 * (r - x.re));
}

if (x.re <= 0) {
  im = 0.5 * Math.sqrt(2.0 * (r - x.re));
}
else {
  im = Math.abs(x.im) / Math.sqrt(2 * (r + x.re));
}

if (x.im >= 0) {
  return new Complex(re, im);
}
else {
  return new Complex(re, -im);
}

Writing this code was harder than I expected. Now, I don’t consider myself to be an outstanding programmer, but I’m reasonably familiar with writing numeric code and felt confident that I understood this particular problem very well. Despite that, I made several mistakes in writing this dozen-line patch.

The first mistake was making a few invalid simplifications in Herbie’s output. Herbie’s output, above, involved squaring \(y\) inside a square root. At first, I just moved \(y\) outside the square root, reasoning that \(\sqrt{x^2} = x\). This is not the case: the sign is wrong for negative \(x\). I also made a mistake in the comparison: the alternate form returns NaN when \(x = 0\), so I had to be careful to send this case to the original expressions.

Luckily, math.js had tests that allowed me to catch both of these bugs before submitting the patch. To make sure the improvement in accuracy would not be removed by future commits, I added a test that specifically tested the improved accuracy. The bug is worst for large, negative \(x\) and small \(y\), in which case the correct real part is approximately \(y / 2 \sqrt{x}\). I used \(10^{10} + i 10^{-10}\) as the test input, with the correct output \(5 · 10^{-16} + i 10^5\). The original code would compute \(0\) for the real part, but the new version did not have this error.

Herbie still has far to go: it would be nice to avoid the tedious and error-prone task of translating between different languages, and Herbie could do more to simplify and optimize the code it produces (because when I do it manually, I make mistakes). Still, in this case and in others, Herbie automatically did a difficult, expert task with minimal human assistance. As Herbie gets better, this should only happen more and more frequently.

I’ll be testing more of math.js in the future, and other open-source mathematics libraries as well. I hope Herbie eventually becomes a tool to tame the difficulty of representing mathematics on computers. Until then, Herbie and I will work on existing numerical code in bits and pieces. I hope we’ll leave the future a little more accurate.

Edit: Thank you James Wilcox, Eric Mullen, and Zach Tatlock for proofreading this post.