Calculate gradient only in a masked area

Question

I have a very large array with only a few small areas of interest. I need to calculate the gradient of this array, but for performance reasons I need this calculation to be restricted to these areas of interest.

I can't do something like this:

phi_grad0[mask] = np.gradient(phi[mask], axis=0)

Because of how fancy indexing works, phi[mask] just becomes a 1D array of the masked pixels, losing spatial information and making the gradient calculation worthless.

np.gradient does handle np.ma.masked_arrays, but the performance is an order of magnitude worse:

import numpy as np
from timeit_context import timeit_context

phi = np.random.randint(low=-100, high=100, size=[100, 100])
phi_mask = np.random.randint(low=0, high=2, size=phi.shape, dtype=np.bool)

with timeit_context('full array'):
    for i2 in range(1000):
        phi_masked_grad1 = np.gradient(phi)

with timeit_context('masked_array'):
    phi_masked = np.ma.masked_array(phi, ~phi_mask)
    for i1 in range(1000):
        phi_masked_grad2 = np.gradient(phi_masked)

This produces the output below:

[full array] finished in 143 ms
[masked_array] finished in 1961 ms

I think its because operations run on masked_arrays are not vectorized, but I'm not sure.

Is there any way of restricting np.gradient so as to achieve better performance?

This timeit_context is a handy timer that works like this, if anyone is interested:

from contextlib import contextmanager
import time

@contextmanager
def timeit_context(name):
    """
    Use it to time a specific code snippet
    Usage: 'with timeit_context('Testcase1'):'
    :param name: Name of the context
    """
    start_time = time.time()
    yield
    elapsed_time = time.time() - start_time
    print('[{}] finished in {} ms'.format(name, int(elapsed_time * 1000)))

Calculate gradient only in a masked area

1 Answers