Bug: _get_weights_and_dims doesn't work properly

[andersy005]

Does anybody have a straightforward example of using the dim keyword on weighted_corr? I'm having issues with it. I'm testing the p-value function on, say, a lat/lon grid of data.

import numpy as np
import xarray as xr
import esmlab

base = np.random.rand(100,10,10)
x = xr.DataArray(np.random.rand(100,10,10) + base, dims=['time','lat','lon'])
y = xr.DataArray(np.random.rand(100,10,10) + base, dims=['time','lat','lon'])
# just looking to correlate over the 'time' dimension so a grid of correlation
# coefficients is returned.
r = esmlab.statistics.weighted_corr(x, y, dim=['time'], weights=None) 

Originally posted by @bradyrx in #90 (comment)

[andersy005]

It appears that the issue is in_apply_nan_mask(). This function assumes that all dimensions should be used:

esmlab/esmlab/statistics.py

Lines 51 to 53 in 479fa70

	weights = xr.DataArray(np.ones(dims_shape), dims=op_over_dims)
	weights = _apply_nan_mask(weights, x, y)

esmlab/esmlab/statistics.py

Lines 21 to 32 in 479fa70

	def _apply_nan_mask(weights, x, y=None):
	# If y is specified, make sure x and y have same shape
	if y is not None and isinstance(y, xr.DataArray):
	assert x.shape == y.shape
	valid = x.notnull() & y.notnull()
	else:
	valid = x.notnull()
	# Apply nan mask
	return weights.where(valid)

As a consequence, the shape of the returned weights is always equal to the input array's shape, even when the user wants to apply the operation along a particular dimension.

To fix this, we will need to come up with a way of doing valid = x.notnull() while taking into consideration dimensions supplied by the user.

[andersy005]

Fixed in #97