A function for fairly efficiently calculating the correlation between two cubes over arbitrary dimensions.

correlate.py

import numpy as np

import iris

def nDimCorr(cube_a, cube_b, corr_dims):
    """ Calculates the n-D correlation cube over the given dimensions
    
    Returns a cube representing the correlation between two 
    cubes along the given dimensions, which are comparable
    in all other dimensions.
    
    Args:
    
    * cube_a, cube_b (cubes):
        Between which the correlation field will be calculated.
    * corr_dims (list of str):
        the cube dimension names 
        over which to calculate correlations.
            
    Returns:
        cube of correlations
    
    """
    
    try:
        
        # construct lists of the dimensions over which contain
        # data to calc correlations with
        # and dimensions of array of correlations
        # and the associated shapes
        res_ind = []
        res_shape = []
        slice_ind = []
        slice_shape = []
        for i, c in enumerate(cube_a.dim_coords):
            if not c.name() in corr_dims:
                res_ind.append(i)
                res_shape.append(len(c.points))
            else:
                slice_ind.append(i)
                slice_shape.append(len(c.points))
        
        # create arrays for use in calculation
        data_a = cube_a.data.copy()
        data_b = cube_b.data.copy()
        
        # reshape data to be data to correlate in 0th dim and
        # other grid points in 1st dim
        # transpose to group the correlation data dims before the
        # grid point dims
        dim_i_len = np.prod(np.array(cube_a.shape)[slice_ind])
        dim_j_len = np.prod(np.array(cube_a.shape)[res_ind])
        flat_a = data_a.transpose(slice_ind+res_ind).reshape(dim_i_len, dim_j_len)
        flat_b = data_b.transpose(slice_ind+res_ind).reshape(dim_i_len, dim_j_len)
        # shape of the transposed grid cube
        trans_iter_shape = np.array(cube_a.shape)[res_ind]
        
        # calc r
        # calc r
        sa = flat_a - np.mean(flat_a, 0)
        sb = flat_b - np.mean(flat_b, 0)
        flat_corrs = np.sum((sa*sb), 0)/np.sqrt(np.sum(sa**2, 0)*np.sum(sb**2, 0))
        
        # reshape the dims and untranspose
        unmap = [n for n, d in sorted(enumerate(res_ind), key=lambda tup: tup[1])]
        corrs = flat_corrs.reshape(trans_iter_shape).transpose(unmap)
        
        # construct cube to hold correlation results
        corrs_cube = iris.cube.Cube(corrs)
        corrs_cube.long_name = "Pearson's r"
        corrs_cube.units = "no_unit"
        corrs_cube.add_history("%s correlation between %s and %s" % (corr_dims, cube_a.name(), cube_b.name()))
        for i, dim in enumerate(res_ind):
            c = cube_a.dim_coords[dim]
            corrs_cube.add_dim_coord(c, i)
            
        return corrs_cube
    
    # cubes are not the same size
    except IndexError:
        print "Cube are incompatible"