Round real and imaginary elements of numpy array with halves rounded toward zero or away from zero

Question

I have a numpy array of complex numbers and need to create a new array with rounded real and imaginary parts, where the rounding at half is either toward zero or away from zero.

There are several recommendations on stackoverflow for using the decimal package which allow one to specify different types of rounding. For an array of complex numbers x, the following code worked, but was very slow:

    rounded_array = np.array([
        float(Decimal(x.real).quantize(0, rounding=ROUND_HALF_DOWN)) + 1j * \
        float(Decimal(x.imag).quantize(0, rounding=ROUND_HALF_DOWNs)) for x in arr])

What are some simple but faster alternatives to this?

The precise meanings of ROUND_HALF_UP and ROUND_HALF_DOWN are shown here: https://docs.python.org/3/library/decimal.html#decimal.ROUND_HALF_UP.

To be very clear, for rounding away from zero or toward zero, for say the real parts of the complex number, I seek (notice the differences at the halves)

      toward zero(ROUND_HALF_DOWN) away from zero (ROUND_HALF_UP)
-4.00         -4.0                 -4.0
-3.75         -4.0                 -4.0
-3.50         -3.0                 -4.0
-3.25         -3.0                 -3.0
-3.00         -3.0                 -3.0
-2.75         -3.0                 -3.0
-2.50         -2.0                 -3.0
-2.25         -2.0                 -2.0
-2.00         -2.0                 -2.0
-1.75         -2.0                 -2.0
-1.50         -1.0                 -2.0
-1.25         -1.0                 -1.0
-1.00         -1.0                 -1.0
-0.75         -1.0                 -1.0
-0.50         -0.0                 -1.0
-0.25         -0.0                 -0.0
 0.00          0.0                  0.0
 0.25          0.0                  0.0
 0.50          0.0                  1.0
 0.75          1.0                  1.0
 1.00          1.0                  1.0
 1.25          1.0                  1.0
 1.50          1.0                  2.0
 1.75          2.0                  2.0
 2.00          2.0                  2.0
 2.25          2.0                  2.0
 2.50          2.0                  3.0
 2.75          3.0                  3.0
 3.00          3.0                  3.0
 3.25          3.0                  3.0
 3.50          3.0                  4.0
 3.75          4.0                  4.0
 4.00          4.0                  4.0

The accepted solution to How to always round up a XX.5 in numpy is both slow and does not provide the type of rounding I'm interested in.

Answer 1

The idea behind the linked answer is sound, but np.vectorize does not really vectorize anything. Let's look at ROUND_HALF_DOWN, which rounds towards zero. This means that for x >= 0, you want to get ceil(x - 0.5). For x < 0 you want to do floor(x + 0.5). This can be accomplished using a vectorized operation, e.g.

mask = (x >= 0)
output = x.copy()
np.add(output, -0.5, where=mask, out=output)
np.add(output, 0.5, where=~mask, out=output)
np.ceil(output, where=mask, out=output)
np.floor(output, where=~mask, out=output)

More space-consumingly, but less verbosely:

mask = (x >= 0)
output = np.empty_like(x)
output[mask] = np.ceil(x[mask] - 0.5)
output[~mask] = np.floor(x[~mask] + 0.5)

To do the operations in-place, just operate on x instead of output. The second part is treating the real and imaginary parts of the array separately. For contiguous arrays, this is easily done with a view, e.g.:

x = x.view(np.float)

You can convert back with

x = x.view(np.complex)

If your arrays are non-contiguous, your best bet is to process the real and imaginary components separately. x.real and x.imag are views into the data, so you can manipulate it in-place.

TL;DR

def round_half_down(arr):
    output = arr.copy().view(np.float)   # This won't fail because the copy is contiguous
    mask = (output >= 0)
    np.subtract(output, 0.5, where=mask, out=output)
    np.ceil(output, where=mask, out=output)
    np.invert(mask, out=mask)
    np.add(output, 0.5, where=mask, out=output)
    np.floor(output, where=mask, out=output)
    return output.view(np.complex)