Calculate Square Root of a Number in Python

Introduction

The square root of a number is a very common mathematical function used in all aspects of science - physics, mathematics, computer science, etc. Square roots of numbers and expressions are very frequent in formulas in all matters of science, and especially in the way we represent reality - by modeling what we can observe with calculus.

In this article, we'll take a look at various ways to calculate a square root of a number in Python. Finally, we'll do a Performance Benchmark with constant and random numbers, as well as lists of random numbers to put all of the approaches to the test.

Calculate Square Root in Python with NumPy

NumPy is a scientific computation library, which found itself present in many applications and use cases. Naturally, it has many wrappers of mathematical functions as helper methods.

If not already installed, you can install it via pip:

$ pip install numpy

In terms of NumPy, the sqrt() function calculates the square root of a number, and returns the result:

import numpy as np
x = np.sqrt(2)
print(x)

This results in:

1.4142135623730951

Aside from taking a single variable as an argument, sqrt() is also able to parse through lists and return a list of square roots:

arr = [2, 3, 5, 7]
roots = np.sqrt(arr)
print(roots)

This results in:

[1.41421356 1.73205081 2.23606798 2.64575131]

The sqrt() function has a limitation though - it cannot calculate a square root of a negative number, because the square root operation with real numbers is only defined for positive numbers.

Attempting to insert -4 into the sqrt() function will result in an exception:

print(np.sqrt(-4))

Trying to calculate a square root of a negative number will result with a warning and a nan value:

RuntimeWarning: invalid value encountered in sqrt
nan

Calculate Square Root of Complex Number with Numpy

Fortunately, NumPy isn't constrained to only work with real numbers - it can work with complex numbers as well:

import numpy as np

complex_number = -1 + 1j
complex_array = [-2, 3, complex_number]

complex_root = np.sqrt(complex_number)
complex_array_roots = np.sqrt(complex_array)

print(f"Square root of '{complex_number}':\n {complex_root}")
print(f"Square roots of '{complex_array}':\n {complex_array_roots}")

If there's at least one complex number in a list, all of the numbers will be cast to and treated as complex, so even negative integers can be added:

Square root of '(-1+1j)':
 (0.45508986056222733+1.09868411346781j)
Square roots of '[-2, 3, (-1+1j)]':
 [0.        +1.41421356j 1.73205081+0.j         0.45508986+1.09868411j]

Python's math Module

The math module is a standard module packaged with Python. It's always available, but has to be imported, and provides wrappers for some common functions, such as the square root, powers, etc:

import math

math.sqrt()

The sqrt() function of the math module is a straightforward function that returns the square root of any positive number:

print(math.sqrt(2))

This results in:

1.4142135623730951

Unlike NumPy's sqrt() function, it can only work on a single element, so if you want to calculate the square root of all elements in a list, you'll have to use a for loop or a list comprehension:

import math

arr = [2, 3, 5, 7]
roots = []

for x in arr:
    roots.append(math.sqrt(x))

# OR
roots = [math.sqrt(x) for x in arr]

In both cases, the roots list will contain:

[1.4142135623730951, 1.7320508075688772, 2.23606797749979, 2.6457513110645907]

math.pow()

A square root of a number can also be calculated by raising a number to a power of ½:

$$
\sqrt x = x^{\frac 1 2}
$$

So really, finding the square root of a number can be expressed as raising the number to a power of ½. math.pow() takes two arguments - the base and the exponent, and rises the base to the power of an exponent:

print(math.pow(2, 0.5))

Naturally, this results in:

1.4142135623730951

The ** Operator

The ** operator is a binary operator which means it works with two values, just like regular multiplication with * does. However, as it is an operator used for exponentiation, we raise its left argument to the power of its right argument.

This approach can be used in the same form as the previous one:

print(2 ** 0.5)

And it also results in:

1.4142135623730951

The pow() Function

Python has another, built-in pow() method that doesn't require an import of the math module. This method is technically different from the math.pow() method internally.

math.pow() implicitly casts elements to doubles, while pow() uses the object's internal implementation, based around the ** operator. While this difference in implementation may warrant the use of one or the other in certain contexts, if you're just calculating the square root of a number, you won't really see the difference:

print(pow(2, 0.5))

This results in:

1.4142135623730951

Performance Benchmark

So which one yields the best performance, and which one should you choose? As usual, there isn't one clear cut winner, and it depends on the usage of the methods. Namely, if you're working with constant numbers, random numbers or an array of random numbers on a larger scale - these methods will perform differently.

Let's test them all out on constant numbers, random numbers and arrays of random numbers:

import timeit

print("Time to execute 100k operations on constant number: \n")
print("math.sqrt(): %ss" % timeit.timeit("math.sqrt(100)", setup="import math", number=100000))
print("math.pow(): %ss" % timeit.timeit("math.pow(100, 0.5)", setup="import math", number=100000))
print("pow(): %ss" % timeit.timeit("pow(100, 0.5)", number=100000))
print("np.sqrt(): %ss" % timeit.timeit("np.sqrt(100)", setup="import numpy as np", number=100000))
print("** operator: %ss" % timeit.timeit("100 ** 0.5", number=100000))

print("\nTime to execute 100k operations on random number: \n")
print("math.sqrt() %ss" % timeit.timeit("math.sqrt(random.random())", setup="import math; import random;", number=100000))
print("math.pow(): %ss" % timeit.timeit("math.pow(random.random(), 0.5)", setup="import math; import random", number=100000))
print("pow(): %ss" % timeit.timeit("pow(random.random(), 0.5)", setup="import random", number=100000))
print("np.sqrt(): %ss" % timeit.timeit("np.sqrt(random.random())", setup="import numpy as np; import random", number=100000))
print("** operator: %ss" % timeit.timeit("random.random() ** 0.5", setup="import random", number=100000))

print("\nTime to execute 100k operations on list of random numbers: \n")
print("math.sqrt() %ss" % timeit.timeit("[math.sqrt(x) for x in np.random.rand(100)]", setup="import math; import numpy as np;", number=100000))
print("math.pow(): %ss" % timeit.timeit("[math.pow(x, 0.5) for x in np.random.rand(100)]", setup="import math; import numpy as np;", number=100000))
print("pow(): %ss" % timeit.timeit("[pow(x, 0.5) for x in np.random.rand(100)]", setup="import numpy as np;", number=100000))
print("np.sqrt(): %ss" % timeit.timeit("np.sqrt(np.random.rand(100))", setup="import numpy as np; import numpy as np;", number=100000))
print("** operator: %ss" % timeit.timeit("np.random.rand(100) ** 0.5", setup="import numpy as np", number=100000))

We've passed all of the methods outlined above through the same test - a constant number (which is likely to be cached for optimization), a random number on each of the 100k iterations, and a list of 100 random numbers.

Running this piece of code results in:

Time to execute 100k operations on constant number: 

math.sqrt(): 0.014326499999999999s
math.pow(): 0.0165132s
pow(): 0.018766599999999994s
np.sqrt(): 0.10575379999999998s
** operator: 0.0006493000000000193s

Time to execute 100k operations on random number: 

math.sqrt() 0.019939999999999958s
math.pow(): 0.022284300000000035s
pow(): 0.0231711s
np.sqrt(): 0.09066460000000004s
** operator: 0.018928s

Time to execute 100k operations on list of random numbers: 

math.sqrt() 2.7786073s
math.pow(): 2.9986906s
pow(): 3.5157339999999992s 
np.sqrt(): 0.2291957s
** operator: 0.2376024000000001s

With constant numbers - the math.pow(), math.sqrt() and pow() functions significantly outperform NumPy's sqrt() function, as they can better utilize caching in the CPU on the language-level.

With random numbers, the caching doesn't work as well and we see smaller discrepancies.

With lists of random numbers, np.sqrt() outperforms all three built-in methods significantly, and the ** operator performs in the same ball-park.

To summarize:

Depending on the concrete input you're dealing with - you'll choose between these functions. While it may seem like they'll all perform well, and while in most cases, it won't make much of a difference, when dealing with huge datasets, even a 10% decrease in processing time can help in the long-run.

Depending on the data you're processing - test the different approaches on your local machine.

Conclusion

In this short article, we've taken a look at several ways to compute the Square Root of a number in Python.

We've taken a look at the math module's pow() and sqrt() functions, as well as the built-in pow() function, NumPy's sqrt() function and the ** operator. Finally, we've benchmarked the methods to compare their performance on different types of input - constant numbers, random numbers and lists of random numbers.