def regress(y, x, axis=0, reverse=False): """ Regress y with x (or x with y if reverse is True) where x is an 1D array where x is 1D of shape (M,), y is (...,M,...) axis specifies where M is in y Args: y (numpy array) x (numpy array) axis (int) reverse (bool): if False (default), then y = slope*x + intercept if True, then x = slope*y + intercept Returns: slope, intercept, p, corr where p is the p-value of the correlation. corr is the correlation coefficient """ if x.squeeze().ndim != 1: raise Exception("x is supposed to be 1D") x = x.squeeze() if y.shape[axis] != len(x): raise Exception("The {}-th axis of y(={}) should match the length"+\ "of x(={})".format(axis, y.shape[axis], len(x))) # Slice object for making x shape broadcast-able to y xtoy = (numpy.newaxis,)*axis + (slice(None),) + \ (numpy.newaxis,)*(y.ndim-axis-1) if reverse: # x regressed with y x_dev = x[xtoy] - x.mean() y_dev = y-keepdims.mean(y, axis) slope = keepdims.sum(x_dev*y_dev, axis)/\ keepdims.sum(numpy.power(y-y.mean(), 2), axis) intercept = (keepdims.mean(x[xtoy], axis)-\ slope*keepdims.mean(y, axis)).squeeze() slope = slope.squeeze() intercept = intercept.squeeze() else: # y regressed with x x_dev = x[xtoy]-x.mean() y_dev = y-keepdims.mean(y, axis) slope = keepdims.sum((x_dev*y_dev), axis)/\ keepdims.sum(numpy.power(x-x.mean(), 2)) intercept = (keepdims.mean(y, axis)-\ slope*keepdims.mean(x[xtoy], axis)).squeeze() slope = slope.squeeze() intercept = intercept.squeeze() def compute_sqrt_x2(x, axis=None): return numpy.ma.sqrt(keepdims.sum(numpy.power(x-x.mean(), 2), axis)) def compute_nd(y, axis=None): if axis is None: nd = len(y) else: nd = y.shape[axis] nd *= numpy.ones([n for idim, n in enumerate(y.shape) if idim != axis]) if isinstance(y, numpy.ma.core.MaskedArray) and y.mask.any(): nd = (~y.mask).sum(axis=axis) return nd sqrt_y2 = compute_sqrt_x2(y, axis=axis) sqrt_x2 = compute_sqrt_x2(x) if reverse: corr = (slope*sqrt_y2/sqrt_x2).squeeze() else: corr = (slope*sqrt_x2/sqrt_y2).squeeze() # calculate degree of freedom if isinstance(x, numpy.ma.core.MaskedArray): if x.mask.any(): nd = compute_nd(x[xtoy]*y, axis=axis) else: nd = compute_nd(y, axis=axis) else: nd = compute_nd(y, axis=axis) # t-test ts = numpy.abs(corr)/numpy.ma.sqrt((1-corr*corr)/(nd-2)) p = 1 - scipy.stats.t.sf(ts, nd) return slope, intercept, p, corr
def compute_sqrt_x2(x, axis=None): return numpy.ma.sqrt(keepdims.sum(numpy.power(x-x.mean(), 2), axis))