Covariance between a 1-d and a 3-d array in Python

Question

I have a dataset created through xarray with its assigned coordinates and dimensions. Obtained from that, I also have two variables: a 1-dimensional array and a 3-dimensional one, with the same coordinate as the first one and two additional ones. I want to obtain the covariance of both in their shared coordinate "memb" for each point in the 2-d space defined by the two coordinates which aren't shared by both and make that a matrix.

In other words, a variable is defined by "memb" and another one is defined by "memb", "north_south" and "west_east". I want to find the memb covariance for each north_south and west_east point and assign it to a variable with a value assigned to each north_south and west_east value.

To obtain it at one point I can run the following code and obtain the desired result:

numpy.cov(var_1,var_2.isel(north_south=1,west_east=1)[0][1]

I want to assign this to a variable which will have the dimensions north_south and west_east. I think I know how to make it work with for blocks, but how can i assign it to a variable with the two dimensions at each point?

Answer 1

The method apply_along_axis seems appropriate. Example:

import numpy as np
a = np.random.uniform(size=(5,))
b = np.random.uniform(size=(5, 3, 2))
c = np.apply_along_axis(lambda x: np.cov(a, x)[0][1], 0, b)

Here c is a 2D array of size 3 by 2. The second parameter of apply_along_axis specifies that the axis of b along which to work is 0th axis (could be another one, as long as it matches the size of 1D array a ). The lambda just computes the covariance, returning the scalar value of interest.