Applying scipy.ndimage.convolve to tridimensional xarray DataArray - python

I have a 3D xarray DataArray with dimensions x, y, z, and I'm trying to apply scipy.ndimage.convolve over each x-y plane while maintaining the output as a DataArray. Naturally I'm trying to use xr.apply_ufunc to do that. If I do it for only one plane it works perfectly:
da=xr.DataArray(np.random.rand(5,5,5), dims=("x", "y", "z"))
kernel=np.ones((3,3))
from scipy.ndimage import convolve
conv1 = lambda x: convolve(x, kernel, mode="wrap")
print(xr.apply_ufunc(conv1, da[:,:,0])) # works successfully
I'm now trying to come up with a way to do the same for every x-y plane. What I thought was going to work was using np.apply_along_axis or np.apply_over_axes, but none of them work.
I could iterate over the axis, put everything in a list, and concatenate, but I'm trying to use xr.apply_ufunc to avoid problems with the attributes. Is there a way to do that?
Here's an example of something that I thought should work, but that doesn't:
np.apply_over_axes(conv1, c, axes=(0,1))
but this fails with
TypeError: <lambda>() takes 1 positional argument but 2 were given

How about using a kernel with shape (3, 3, 1) instead (3, 3)?
kernel2d = np.ones((3, 3))
conv2d = lambda x: convolve(x, kernel2d, mode="wrap")
result2d = xr.apply_ufunc(conv2d, da[:, :, 0])
kernel3d = np.ones((3, 3, 1))
conv3d = lambda x: convolve(x, kernel3d, mode="wrap")
result3d = xr.apply_ufunc(conv3d, da)
(result2d == result3d[:, :, 0]).all() # -> True
Another option is to use vectorization logic in xr.apply_ufunc, which may be closer to what you tried to do
kernel = np.ones((3, 3))
conv = lambda x: convolve(x, kernel, mode="wrap")
result = xr.apply_ufunc(conv, da, input_core_dims=[['x', 'y']],
output_core_dims=[['x', 'y']],
vectorize=True)
(result2d == result.transpose('x', 'y', 'z')).all() # --> True
This option is only prepared for convenience and therefore it might be much slower than the first one where the calculation is vectorized.

A possible answer that I came up with is to manually do this:
def conv_rx(da, axis="z"):
planes = [ xr.apply_ufunc(conv1, da.sel(z=z)) for z in da.z ]
new = xr.concat(planes, dim=axis)
return new.transpose(*da.dims)
which yields the correct result. However, I'm not very happy with this since it's not elegant and it's pretty slow.

Related

2d convolution gives not the desired output

I want to use the 2D convolution in the same way I did here in 1D. Unfortunately the output in the former case does not have the desired shape. Let n = 5, then
h_0 = (1 / 4) * np.array([1, 2, 1])
x = np.random.rand(n)
np.convolve(h_0, x, 'same')
>>> array([0.65498075, 0.72729356, 0.51417706, 0.34597679, 0.1793755])
but
h_00 = np.kron(h_0, h_0)
h_00 = np.reshape(h_00, (3, 3))
x = np.random.rand(n, n)
scipy.signal.convolve2d(h_00, x, 'same', boundary='symm')
>>> array([[1.90147294, 1.6541233 , 1.82704077],
[1.55228912, 1.3641027 , 1.55536069],
[1.61190909, 1.45159935, 1.58266083]])
I would have expected a (5, 5) output array.
The docs for scipy.signal.convolve2d regarding the mode parameter clearly state
mode
...
same
    The output is the same size as in1, centered with respect to the ‘full’ output.
So, given that you pass the kernel first, your output will be the same size as the kernel, not the array you are filtering. To fix, swap the first two inputs:
scipy.signal.convolve2d(x, h_00, 'same', boundary='symm')
Confusion likely arises from the behavior of numpy.convolve, which does the following:
mode : {‘full’, ‘valid’, ‘same’}, optional
...
‘same’:
    Mode ‘same’ returns output of length max(M, N). Boundary effects are still visible.
Numpy interprets the larger array as the kernel regardless of argument order. This is possible because with a single dimension, there is always an unambiguous winner.

Vector dot product along one dimension for multidimensional arrays

I want to compute the sum product along one dimension of two multidimensional arrays, using Theano.
I'll describe precisely what I want to do using numpy first. numpy.tensordot and numpy.dot seem to always do a matrix product, whereas I'm in essence looking for a batched equivalent of a vector product. Given x and y, I want to compute z like so:
x = np.random.normal(size=(200, 2, 2, 1000))
y = np.random.normal(size=(200, 2, 2))
# this is how I now approach it:
z = np.sum(y[:,:,:,np.newaxis] * x, axis=1)
# z is of shape (200, 2, 1000)
Now I know that numpy.einsum would probably be able to help me here, but again, I want to do this particular computation in Theano, which does not have an einsum equivalent. I will need to use dot, tensordot, or Theano's specialized einsum subset functions batched_dot or batched_tensordot.
The reason I'm looking to change my approach to this is performance; I suspect that using builtin (CUDA) dot products will be faster than relying on broadcasting, element-wise product, and sum.
In Theano, none of the dimensions of three and four dimensional tensors are broadcastable. You have to explicitly set them. Then the Numpy principles will work just fine. One way to do this is to use T.patternbroadcast. To read more about broadcasting, refer this.
You have three dimensions in one of the tensors. So first you need to append a singleton dimension at the end and then make that dimension broadcastable. These two things can be achieved with a single command - T.shape_padaxis. The entire code is as follows:
import theano
from theano import tensor as T
import numpy as np
X = T.ftensor4('X')
Y = T.ftensor3('Y')
Y_broadcast = T.shape_padaxis(Y, axis=-1) # appending extra dimension and making it
# broadcastable
Z = T.sum((X*Y_broadcast), axis=1) # element-wise multiplication
f = theano.function([X, Y], Z, allow_input_downcast=True)
# Making sure that it works and gives correct results
x = np.random.normal(size=(3, 2, 2, 4))
y = np.random.normal(size=(3, 2, 2))
theano_result = f(x,y)
numpy_result = np.sum(y[:,:,:,np.newaxis] * x, axis=1)
print np.amax(theano_result - numpy_result) # prints 2.7e-7 on my system, close enough!
I hope this helps.

Dealing with dimension collapse in python arrays

A recurring error I run into when using NumPy is that an attempt to index an array fails because one of the dimensions of the array was a singleton, and thus that dimension got wiped out and can't be indexed. This is especially problematic in functions designed to operate on arrays of arbitrary size. I'm looking for the cheapest, most universal way to avoid this error.
Here's an example:
import numpy as np
f = (lambda t, u, i=0: t[:,i]*u[::-1])
a = np.eye(3)
b = np.array([1,2,3])
f(a,b)
f(a[:,0],b[1])
The first call works as expected. The second call fails in two ways: 1) t can't be indexed by [:,0] because is has shape (3,), and 2) u can't be indexed at all because it's a scalar.
Here are the fixes that occur to me:
1) Use np.atleast_1d and np.atleast_2d etc. (possibly with conditionals to make sure that the dimensions are in the right order) inside f to make sure that all parameters have the dimensions they need. This precludes use of lambdas, and can take a few lines that I would rather not need.
2) Instead of writing f(a[:,0],b[1]) above, use f(a[:,[0]],b[[1]]). This is fine, but I always have to remember to put in the extra brackets, and if the index is stored in a variable you might not know if you should put the extra brackets in or not. E.g.:
idx = 1
f(a[:,[0]],b[[idx]])
idx = [2,0,1]
f(a[:,[0]],b[idx])
In this case, you would seem to have to call np.atleast_1d on idx first, which may be even more cumbersome than putting np.atleast_1d in the function.
3) In some cases I can get away with just not putting in an index. E.g.:
f = lambda t, u: t[0]*u
f(a,b)
f(a[:,0],b[0])
This works, and is apparently the slickest solution when it applies. But it doesn't help in every case (in particular, your dimensions have to be in the right order to begin with).
So, are there better approaches than the above?
There are lots of ways to avoid this behaviour.
First, whenever you index into a dimension of an np.ndarray with a slice rather than an integer, the number of dimensions of the output will be the same as that of the input:
import numpy as np
x = np.arange(12).reshape(3, 4)
print x[:, 0].shape # integer indexing
# (3,)
print x[:, 0:1].shape # slice
# (3, 1)
This is my preferred way of avoiding the problem, since it generalizes very easily from single-element to multi-element selections (e.g. x[:, i:i+1] vs x[:, i:i+n]).
As you've already touched on, you can also avoid dimension loss by using any sequence of integers to index into a dimension:
print x[:, [0]].shape # list
# (3, 1)
print x[:, (0,)].shape # tuple
# (3, 1)
print x[:, np.array((0,))].shape # array
# (3, 1)
If you choose to stick with integer indices, you can always insert a new singleton dimension using np.newaxis (or equivalently, None):
print x[:, 0][:, np.newaxis]
# (3, 1)
print x[:, 0][:, None]
# (3, 1)
Or else you could manually reshape it to the correct size (here using -1 to infer the size of the first dimension automatically):
print x[:, 0].reshape(-1, 1).shape
# (3, 1)
Finally, you can use an np.matrix rather than an np.ndarray. np.matrix behaves more like a MATLAB matrix, where singleton dimensions are left in whenever you index with an integer:
y = np.matrix(x)
print y[:, 0].shape
# (3, 1)
However, you should be aware that there are a number of other important differences between np.matrix and np.ndarray, for example the * operator performs elementwise multiplication on arrays, but matrix multiplication on matrices. In most circumstances it's best to stick to np.ndarrays.

Why does theano conv2d add empty dimension?

I am playing around with some simple Theano code, and I ran into the following:
import numpy
import theano
from theano import tensor
from theano.tensor.signal.conv import conv2d
m = tensor.fmatrix()
w = numpy.ones([10,1], dtype=numpy.float32)
c = conv2d(m,w)
f = theano.function([m], c)
print f(numpy.ones([100,100], dtype=numpy.float32)).shape
Result:
(1, 91, 100)
The result of a 2d convolution of 2d inputs is expected to be 2d, but it is actually 3d. Why?
The docstring of conv2d says signal.conv.conv2d performs a basic 2D convolution of the input with the
given filters. (note the plural)
You could pass it several filters and it will return the convolutions with all of those. Try e.g.
c = conv2d(m,np.array([w, w, w]))
f = theano.function([m], c)
print f(numpy.ones([100,100], dtype=numpy.float32)).shape # outputs (3, 91, 100)
So it seems that by default it will add a degenerate axis if you only pass 1 filter (probably because it adds this axis internally to your filter if you didn't pass it in that way yourself. In other words, it doesn't keep track of the input shape in order to return something that corresponds. Looks like a design choice more than anything else.)

numpy broadcast from first dimension

In NumPy, is there an easy way to broadcast two arrays of dimensions e.g. (x,y) and (x,y,z)? NumPy broadcasting typically matches dimensions from the last dimension, so usual broadcasting will not work (it would require the first array to have dimension (y,z)).
Background: I'm working with images, some of which are RGB (shape (h,w,3)) and some of which are grayscale (shape (h,w)). I generate alpha masks of shape (h,w), and I want to apply the mask to the image via mask * im. This doesn't work because of the above-mentioned problem, so I end up having to do e.g.
mask = mask.reshape(mask.shape + (1,) * (len(im.shape) - len(mask.shape)))
which is ugly. Other parts of the code do operations with vectors and matrices, which also run into the same issue: it fails trying to execute m + v where m has shape (x,y) and v has shape (x,). It's possible to use e.g. atleast_3d, but then I have to remember how many dimensions I actually wanted.
how about use transpose:
(a.T + c.T).T
numpy functions often have blocks of code that check dimensions, reshape arrays into compatible shapes, all before getting down to the core business of adding or multiplying. They may reshape the output to match the inputs. So there is nothing wrong with rolling your own that do similar manipulations.
Don't offhand dismiss the idea of rotating the variable 3 dimension to the start of the dimensions. Doing so takes advantage of the fact that numpy automatically adds dimensions at the start.
For element by element multiplication, einsum is quite powerful.
np.einsum('ij...,ij...->ij...',im,mask)
will handle cases where im and mask are any mix of 2 or 3 dimensions (assuming the 1st 2 are always compatible. Unfortunately this does not generalize to addition or other operations.
A while back I simulated einsum with a pure Python version. For that I used np.lib.stride_tricks.as_strided and np.nditer. Look into those functions if you want more power in mixing and matching dimensions.
as another angle: if you encounter this pattern frequently, it may be useful to create a utility function to enforce right-broadcasting:
def right_broadcasting(arr, target):
return arr.reshape(arr.shape + (1,) * (target.ndim - arr.ndim))
Although if there are only two types of input (already having 3 dims or having only 2), id say the single if statement is preferable.
Indexing with np.newaxis creates a new axis in that place. Ie
xyz = #some 3d array
xy = #some 2d array
xyz_sum = xyz + xy[:,:,np.newaxis]
or
xyz_sum = xyz + xy[:,:,None]
Indexing in this way creates an axis with shape 1 and stride 0 in this location.
Why not just decorate-process-undecorate:
def flipflop(func):
def wrapper(a, mask):
if len(a.shape) == 3:
mask = mask[..., None]
b = func(a, mask)
return np.squeeze(b)
return wrapper
#flipflop
def f(x, mask):
return x * mask
Then
>>> N = 12
>>> gs = np.random.random((N, N))
>>> rgb = np.random.random((N, N, 3))
>>>
>>> mask = np.ones((N, N))
>>>
>>> f(gs, mask).shape
(12, 12)
>>> f(rgb, mask).shape
(12, 12, 3)
Easy, you just add a singleton dimension at the end of the smaller array. For example, if xyz_array has shape (x,y,z) and xy_array has shape (x,y), you can do
xyz_array + np.expand_dims(xy_array, xy_array.ndim)

Categories

Resources