calculating windowed probability mass function on an image using numpy and scipy

Question

I start with a 2D image array in numpy, for which I want to perform some windowed processing. I'm looking for a vectorized approach, that for each window, first calculates a probability mass function, and then replaces the center pixel with an operation on that mass function. Further, it needs to work with nonlinear operations, like log(x). For example: given the following 3x3 window:

[[1, 1, 1],
 [2, 2, 2],
 [3, 4, 5]]

The probability mass function would be something like: p(1) = 0.33, p(2) = 0.33, p(3) = p(4) = p(5) = 0.11. I started with this question yesterday, which led me to the following process, giving me probabilities for each value in a window ksize (3x3 in the example above), relative to the other values in that window.

def process(a, ksize=3):
  # get the number of times a value occurs in each kernel window
  b = windowed_occurences(a, ksize)

  # convert to probabilities
  p = a / (ksize * ksize)

def windowed_occurences(a, ksize):
  window_shape = (ksize, ksize)
  d = math.floor(ksize / 2)
  a = np.pad(a, (d, d), 'constant', constant_values=(0, 0))

  # get the windowed array of shape (a.height, a.width, ksize, ksize)
  b = skimage.util.shape.view_as_windows(a, window_shape)

  # replace each element of ksize x ksize kernel with the number of occurances
  # of that element in the kernel
  # process from https://stackoverflow.com/questions/47130141/ @Divakar
  ar2D = b.reshape(-1, b.shape[-2] * b.shape[-1])
  c = bincount2D_vectorized(ar2D)
  d = c[np.arange(ar2D.shape[0])[:, None], ar2D].reshape(b.shape)

  return d

# bincount2D_vectorized from https://stackoverflow.com/a/46256361/ @Divakar
def bincount2D_vectorized(a):    
  N = a.max()+1
  a_offs = a + np.arange(a.shape[0])[:,None]*N
  return np.bincount(a_offs.ravel(), minlength=a.shape[0]*N).reshape(-1,N)

For the example window, this gives:

[[0.33, 0.33, 0.33]
 [0.33, 0.33, 0.33]
 [0.11, 0.11, 0.11]]

Now if I wanted the sum of the logarithms of those probabilities, I could do

s = (np.log(p)).sum((-2, -1))

Gives me s = -13.18 (6*log(1/3)+3*log(1/9)) for my example, which is incorrect. This would be correct if each value in the window were unique (p(x) = 0.11 across the window), but is incorrect in the example above, since there are duplicate elements. I need a way to either zero the duplicate values after computing their probabilities (some clever application of np.unique), or some other way to run the calculation. I can't simply sum the probabilities and normalize, since log is nonlinear. A correct result would be s = -8.79 (2*log(1/3)+3*log(1/9)).

Edit: Solution I posted a solution below. It's about 10x faster for my sample images than using a python loopy approach. I'd like an even faster approach, but I doubt that it's possible, since log is nonlinear, the "convolution" is not separable.

Answer 1

After taking a long, and close look at @Divakar's bincount2D_vectorized, I realized the probability mass function was already in there, and just needed to be reshaped. To replace each pixel in the original image a, with a sum over the its kernel of log(p(x)) where p(x) is the probability of finding the pixel value x in that kernel, we can do:

def process(a, ksize=3):
  # get the number of times a value occurs in each kernel window
  window_shape = (ksize, ksize)
  d = math.floor(ksize / 2)
  a = np.pad(a, (d, d), 'constant', constant_values=(0, 0))

  # get the windowed array of shape (a.height, a.width, ksize, ksize)
  b = skimage.util.shape.view_as_windows(a, window_shape)

  # process from https://stackoverflow.com/questions/47130141/ @Divakar
  ar2D = b.reshape(-1, b.shape[-2] * b.shape[-1])
  c = bincount2D_vectorized(ar2D)

  # convert to probabilities
  p = c / (ksize * ksize)
  p[p == 0] = 1 # divide by 0 error

  # sum and reshape
  return np.sum(np.log(p), axis=1).reshape(a.shape)