def _vertical_interpolate(cube, src_levels, levels, interpolation, extrapolation): """Perform vertical interpolation.""" # Determine the source levels and axis for vertical interpolation. z_axis, = cube.coord_dims(cube.coord(axis='z', dim_coords=True)) # Broadcast the 1d source cube vertical coordinate to fully # describe the spatial extent that will be interpolated. src_levels_broadcast = broadcast_to_shape(src_levels.points, cube.shape, cube.coord_dims(src_levels)) # force mask onto data as nan's if np.ma.is_masked(cube.data): cube.data[cube.data.mask] = np.nan # Now perform the actual vertical interpolation. new_data = stratify.interpolate(levels, src_levels_broadcast, cube.data, axis=z_axis, interpolation=interpolation, extrapolation=extrapolation) # Calculate the mask based on the any NaN values in the interpolated data. mask = np.isnan(new_data) if np.any(mask): # Ensure that the data is masked appropriately. new_data = np.ma.array(new_data, mask=mask, fill_value=_MDI) # Construct the resulting cube with the interpolated data. return _create_cube(cube, new_data, src_levels, levels.astype(float))
def test_masked_degenerate(self): # masked arrays can have degenerate masks too a = np.random.random([2, 3]) m = ma.array(a) b = broadcast_to_shape(m, (5, 3, 4, 2), (3, 1)) for i in xrange(5): for j in xrange(4): self.assertMaskedArrayEqual(b[i, :, j, :].T, m)
def test_added_dimensions_transpose(self): # adding dimensions and having the dimensions of the input # transposed a = np.random.random([2, 3]) b = broadcast_to_shape(a, (5, 3, 4, 2), (3, 1)) for i in xrange(5): for j in xrange(4): self.assertArrayEqual(b[i, :, j, :].T, a)
def test_masked(self): # masked arrays are also accepted a = np.random.random([2, 3]) m = ma.array(a, mask=[[0, 1, 0], [0, 1, 1]]) b = broadcast_to_shape(m, (5, 3, 4, 2), (3, 1)) for i in xrange(5): for j in xrange(4): self.assertMaskedArrayEqual(b[i, :, j, :].T, m)
def test_added_dimensions(self): # adding two dimensions, on at the front and one in the middle of # the existing dimensions a = np.random.random([2, 3]) b = broadcast_to_shape(a, (5, 2, 4, 3), (1, 3)) for i in xrange(5): for j in xrange(4): self.assertArrayEqual(b[i, :, j, :], a)
def _compute_mean(data, weights): mapping = set( data.coord_dims('latitude') + data.coord_dims('longitude')) with warnings.catch_warnings(): warnings.simplefilter("ignore") return data.collapsed( ('latitude', 'longitude'), iris.analysis.MEAN, weights=broadcast_to_shape(weights, data.shape, mapping))
def coord_to_cube(cube, coord): """ Convert coordinate points to a cube of the same dimension as the given cube. Parameters ---------- cube: iris.cube.Cube Cube containing the coordinate to be broadcast. coord: str or iris.coords.Coord Coordinate to be broadcast Returns ------- iris.cube.Cube Cube of broadcast coordinate """ if isinstance(coord, str): _coord = cube.coord(coord) else: _coord = coord dim_map = cube.coord_dims(_coord.name()) _data = _coord.points if len(dim_map) > 0: _data = broadcast_to_shape(_data, cube.shape, dim_map) dc = [(c.copy(), cube.coord_dims(c)) for c in cube.dim_coords] ac = [(c.copy(), cube.coord_dims(c)) for c in cube.aux_coords] new_cube = iris.cube.Cube( data=_data, units=_coord.units, long_name=_coord.name(), dim_coords_and_dims=dc, aux_coords_and_dims=ac, ) else: new_cube = iris.cube.Cube(data=_data, standard_name=_coord.name(), units=_coord.units) return new_cube
def test_same_shape(self): # broadcast to current shape should result in no change a = np.random.random([2, 3]) b = broadcast_to_shape(a, a.shape, (0, 1)) self.assertArrayEqual(b, a)
def pearsonr(cube_a, cube_b, corr_coords=None, weights=None, mdtol=1.): """ Calculate the Pearson's r correlation coefficient over specified dimensions. Args: * cube_a, cube_b (cubes): Cubes between which the correlation will be calculated. The cubes should either be the same shape and have the same dimension coordinates or one cube should be broadcastable to the other. * corr_coords (str or list of str): The cube coordinate name(s) over which to calculate correlations. If no names are provided then correlation will be calculated over all common cube dimensions. * weights (numpy.ndarray, optional): Weights array of same shape as (the smaller of) cube_a and cube_b. Note that latitude/longitude area weights can be calculated using :func:`iris.analysis.cartography.area_weights`. * mdtol (float, optional): Tolerance of missing data. The missing data fraction is calculated based on the number of grid cells masked in both cube_a and cube_b. If this fraction exceed mdtol, the returned value in the corresponding cell is masked. mdtol=0 means no missing data is tolerated while mdtol=1 means the resulting element will be masked if and only if all contributing elements are masked in cube_a or cube_b. Defaults to 1. Returns: A cube of the correlation between the two input cubes along the specified dimensions, at each point in the remaining dimensions of the cubes. For example providing two time/altitude/latitude/longitude cubes and corr_coords of 'latitude' and 'longitude' will result in a time/altitude cube describing the latitude/longitude (i.e. pattern) correlation at each time/altitude point. Reference: http://www.statsoft.com/textbook/glosp.html#Pearson%20Correlation """ # Assign larger cube to cube_1 if cube_b.ndim > cube_a.ndim: cube_1 = cube_b cube_2 = cube_a else: cube_1 = cube_a cube_2 = cube_b dim_coords_1 = [coord.name() for coord in cube_1.dim_coords] dim_coords_2 = [coord.name() for coord in cube_2.dim_coords] common_dim_coords = list(set(dim_coords_1) & set(dim_coords_2)) # If no coords passed then set to all common dimcoords of cubes. if corr_coords is None: corr_coords = common_dim_coords # Broadcast weights to shape of cube_1 if necessary. if weights is None or cube_1.shape == cube_2.shape: weights_1 = weights weights_2 = weights else: if weights.shape != cube_2.shape: raise ValueError("weights array should have dimensions {}".format( cube_2.shape)) dims_1_common = [ i for i in range(cube_1.ndim) if dim_coords_1[i] in common_dim_coords ] weights_1 = broadcast_to_shape(weights, cube_1.shape, dims_1_common) weights_2 = weights # Calculate correlations. s1 = cube_1 - cube_1.collapsed( corr_coords, iris.analysis.MEAN, weights=weights_1) s2 = cube_2 - cube_2.collapsed( corr_coords, iris.analysis.MEAN, weights=weights_2) covar = (s1 * s2).collapsed(corr_coords, iris.analysis.SUM, weights=weights_1, mdtol=mdtol) var_1 = (s1**2).collapsed(corr_coords, iris.analysis.SUM, weights=weights_1) var_2 = (s2**2).collapsed(corr_coords, iris.analysis.SUM, weights=weights_2) denom = iris.analysis.maths.apply_ufunc(np.sqrt, var_1 * var_2, new_unit=covar.units) corr_cube = covar / denom corr_cube.rename("Pearson's r") return corr_cube
def pearsonr(cube_a, cube_b, corr_coords=None, weights=None, mdtol=1., common_mask=False): """ Calculate the Pearson's r correlation coefficient over specified dimensions. Args: * cube_a, cube_b (cubes): Cubes between which the correlation will be calculated. The cubes should either be the same shape and have the same dimension coordinates or one cube should be broadcastable to the other. * corr_coords (str or list of str): The cube coordinate name(s) over which to calculate correlations. If no names are provided then correlation will be calculated over all common cube dimensions. * weights (numpy.ndarray, optional): Weights array of same shape as (the smaller of) cube_a and cube_b. Note that latitude/longitude area weights can be calculated using :func:`iris.analysis.cartography.area_weights`. * mdtol (float, optional): Tolerance of missing data. The missing data fraction is calculated based on the number of grid cells masked in both cube_a and cube_b. If this fraction exceed mdtol, the returned value in the corresponding cell is masked. mdtol=0 means no missing data is tolerated while mdtol=1 means the resulting element will be masked if and only if all contributing elements are masked in cube_a or cube_b. Defaults to 1. * common_mask (bool): If True, applies a common mask to cube_a and cube_b so only cells which are unmasked in both cubes contribute to the calculation. If False, the variance for each cube is calculated from all available cells. Defaults to False. Returns: A cube of the correlation between the two input cubes along the specified dimensions, at each point in the remaining dimensions of the cubes. For example providing two time/altitude/latitude/longitude cubes and corr_coords of 'latitude' and 'longitude' will result in a time/altitude cube describing the latitude/longitude (i.e. pattern) correlation at each time/altitude point. Reference: http://www.statsoft.com/textbook/glosp.html#Pearson%20Correlation """ # Assign larger cube to cube_1 if cube_b.ndim > cube_a.ndim: cube_1 = cube_b cube_2 = cube_a else: cube_1 = cube_a cube_2 = cube_b dim_coords_1 = [coord.name() for coord in cube_1.dim_coords] dim_coords_2 = [coord.name() for coord in cube_2.dim_coords] common_dim_coords = list(set(dim_coords_1) & set(dim_coords_2)) # If no coords passed then set to all common dimcoords of cubes. if corr_coords is None: corr_coords = common_dim_coords smaller_shape = cube_2.shape # Match up data masks if required. if common_mask: # Create a cube of 1's with a common mask. if ma.is_masked(cube_2.data): mask_cube = _ones_like(cube_2) else: mask_cube = 1. if ma.is_masked(cube_1.data): # Take a slice to avoid unnecessary broadcasting of cube_2. slice_coords = [dim_coords_1[i] for i in range(cube_1.ndim) if dim_coords_1[i] not in common_dim_coords and np.array_equal(cube_1.data.mask.any(axis=i), cube_1.data.mask.all(axis=i))] cube_1_slice = next(cube_1.slices_over(slice_coords)) mask_cube = _ones_like(cube_1_slice) * mask_cube # Apply common mask to data. if isinstance(mask_cube, iris.cube.Cube): cube_1 = cube_1 * mask_cube cube_2 = mask_cube * cube_2 dim_coords_2 = [coord.name() for coord in cube_2.dim_coords] # Broadcast weights to shape of cubes if necessary. if weights is None or cube_1.shape == smaller_shape: weights_1 = weights weights_2 = weights else: if weights.shape != smaller_shape: raise ValueError("weights array should have dimensions {}". format(smaller_shape)) dims_1_common = [i for i in range(cube_1.ndim) if dim_coords_1[i] in common_dim_coords] weights_1 = broadcast_to_shape(weights, cube_1.shape, dims_1_common) if cube_2.shape != smaller_shape: dims_2_common = [i for i in range(cube_2.ndim) if dim_coords_2[i] in common_dim_coords] weights_2 = broadcast_to_shape(weights, cube_2.shape, dims_2_common) else: weights_2 = weights # Calculate correlations. s1 = cube_1 - cube_1.collapsed(corr_coords, iris.analysis.MEAN, weights=weights_1) s2 = cube_2 - cube_2.collapsed(corr_coords, iris.analysis.MEAN, weights=weights_2) covar = (s1*s2).collapsed(corr_coords, iris.analysis.SUM, weights=weights_1, mdtol=mdtol) var_1 = (s1**2).collapsed(corr_coords, iris.analysis.SUM, weights=weights_1) var_2 = (s2**2).collapsed(corr_coords, iris.analysis.SUM, weights=weights_2) denom = iris.analysis.maths.apply_ufunc(np.sqrt, var_1*var_2, new_unit=covar.units) corr_cube = covar / denom corr_cube.rename("Pearson's r") return corr_cube
def pearsonr(cube_a, cube_b, corr_coords=None, weights=None, mdtol=1., common_mask=False): """ Calculate the Pearson's r correlation coefficient over specified dimensions. Args: * cube_a, cube_b (cubes): Cubes between which the correlation will be calculated. The cubes should either be the same shape and have the same dimension coordinates or one cube should be broadcastable to the other. * corr_coords (str or list of str): The cube coordinate name(s) over which to calculate correlations. If no names are provided then correlation will be calculated over all common cube dimensions. * weights (numpy.ndarray, optional): Weights array of same shape as (the smaller of) cube_a and cube_b. Note that latitude/longitude area weights can be calculated using :func:`iris.analysis.cartography.area_weights`. * mdtol (float, optional): Tolerance of missing data. The missing data fraction is calculated based on the number of grid cells masked in both cube_a and cube_b. If this fraction exceed mdtol, the returned value in the corresponding cell is masked. mdtol=0 means no missing data is tolerated while mdtol=1 means the resulting element will be masked if and only if all contributing elements are masked in cube_a or cube_b. Defaults to 1. * common_mask (bool): If True, applies a common mask to cube_a and cube_b so only cells which are unmasked in both cubes contribute to the calculation. If False, the variance for each cube is calculated from all available cells. Defaults to False. Returns: A cube of the correlation between the two input cubes along the specified dimensions, at each point in the remaining dimensions of the cubes. For example providing two time/altitude/latitude/longitude cubes and corr_coords of 'latitude' and 'longitude' will result in a time/altitude cube describing the latitude/longitude (i.e. pattern) correlation at each time/altitude point. Reference: http://www.statsoft.com/textbook/glosp.html#Pearson%20Correlation This operation is non-lazy. """ # Assign larger cube to cube_1 if cube_b.ndim > cube_a.ndim: cube_1 = cube_b cube_2 = cube_a else: cube_1 = cube_a cube_2 = cube_b dim_coords_1 = [coord.name() for coord in cube_1.dim_coords] dim_coords_2 = [coord.name() for coord in cube_2.dim_coords] common_dim_coords = list(set(dim_coords_1) & set(dim_coords_2)) # If no coords passed then set to all common dimcoords of cubes. if corr_coords is None: corr_coords = common_dim_coords smaller_shape = cube_2.shape # Match up data masks if required. if common_mask: # Create a cube of 1's with a common mask. if ma.is_masked(cube_2.data): mask_cube = _ones_like(cube_2) else: mask_cube = 1. if ma.is_masked(cube_1.data): # Take a slice to avoid unnecessary broadcasting of cube_2. slice_coords = [ dim_coords_1[i] for i in range(cube_1.ndim) if dim_coords_1[i] not in common_dim_coords and np.array_equal( cube_1.data.mask.any(axis=i), cube_1.data.mask.all(axis=i)) ] cube_1_slice = next(cube_1.slices_over(slice_coords)) mask_cube = _ones_like(cube_1_slice) * mask_cube # Apply common mask to data. if isinstance(mask_cube, iris.cube.Cube): cube_1 = cube_1 * mask_cube cube_2 = mask_cube * cube_2 dim_coords_2 = [coord.name() for coord in cube_2.dim_coords] # Broadcast weights to shape of cubes if necessary. if weights is None or cube_1.shape == smaller_shape: weights_1 = weights weights_2 = weights else: if weights.shape != smaller_shape: raise ValueError("weights array should have dimensions {}".format( smaller_shape)) dims_1_common = [ i for i in range(cube_1.ndim) if dim_coords_1[i] in common_dim_coords ] weights_1 = broadcast_to_shape(weights, cube_1.shape, dims_1_common) if cube_2.shape != smaller_shape: dims_2_common = [ i for i in range(cube_2.ndim) if dim_coords_2[i] in common_dim_coords ] weights_2 = broadcast_to_shape(weights, cube_2.shape, dims_2_common) else: weights_2 = weights # Calculate correlations. s1 = cube_1 - cube_1.collapsed( corr_coords, iris.analysis.MEAN, weights=weights_1) s2 = cube_2 - cube_2.collapsed( corr_coords, iris.analysis.MEAN, weights=weights_2) covar = (s1 * s2).collapsed(corr_coords, iris.analysis.SUM, weights=weights_1, mdtol=mdtol) var_1 = (s1**2).collapsed(corr_coords, iris.analysis.SUM, weights=weights_1) var_2 = (s2**2).collapsed(corr_coords, iris.analysis.SUM, weights=weights_2) denom = iris.analysis.maths.apply_ufunc(np.sqrt, var_1 * var_2, new_unit=covar.units) corr_cube = covar / denom corr_cube.rename("Pearson's r") return corr_cube
def pearsonr(cube_a, cube_b, corr_coords=None, weights=None, mdtol=1.): """ Calculate the Pearson's r correlation coefficient over specified dimensions. Args: * cube_a, cube_b (cubes): Cubes between which the correlation will be calculated. The cubes should either be the same shape and have the same dimension coordinates or one cube should be broadcastable to the other. * corr_coords (str or list of str): The cube coordinate name(s) over which to calculate correlations. If no names are provided then correlation will be calculated over all common cube dimensions. * weights (numpy.ndarray, optional): Weights array of same shape as (the smaller of) cube_a and cube_b. Note that latitude/longitude area weights can be calculated using :func:`iris.analysis.cartography.area_weights`. * mdtol (float, optional): Tolerance of missing data. The missing data fraction is calculated based on the number of grid cells masked in both cube_a and cube_b. If this fraction exceed mdtol, the returned value in the corresponding cell is masked. mdtol=0 means no missing data is tolerated while mdtol=1 means the resulting element will be masked if and only if all contributing elements are masked in cube_a or cube_b. Defaults to 1. Returns: A cube of the correlation between the two input cubes along the specified dimensions, at each point in the remaining dimensions of the cubes. For example providing two time/altitude/latitude/longitude cubes and corr_coords of 'latitude' and 'longitude' will result in a time/altitude cube describing the latitude/longitude (i.e. pattern) correlation at each time/altitude point. Reference: http://www.statsoft.com/textbook/glosp.html#Pearson%20Correlation """ # Assign larger cube to cube_1 if cube_b.ndim > cube_a.ndim: cube_1 = cube_b cube_2 = cube_a else: cube_1 = cube_a cube_2 = cube_b dim_coords_1 = [coord.name() for coord in cube_1.dim_coords] dim_coords_2 = [coord.name() for coord in cube_2.dim_coords] common_dim_coords = list(set(dim_coords_1) & set(dim_coords_2)) # If no coords passed then set to all common dimcoords of cubes. if corr_coords is None: corr_coords = common_dim_coords # Broadcast weights to shape of cube_1 if necessary. if weights is None or cube_1.shape == cube_2.shape: weights_1 = weights weights_2 = weights else: if weights.shape != cube_2.shape: raise ValueError("weights array should have dimensions {}". format(cube_2.shape)) dims_1_common = [i for i in range(cube_1.ndim) if dim_coords_1[i] in common_dim_coords] weights_1 = broadcast_to_shape(weights, cube_1.shape, dims_1_common) weights_2 = weights # Calculate correlations. s1 = cube_1 - cube_1.collapsed(corr_coords, iris.analysis.MEAN, weights=weights_1) s2 = cube_2 - cube_2.collapsed(corr_coords, iris.analysis.MEAN, weights=weights_2) covar = (s1*s2).collapsed(corr_coords, iris.analysis.SUM, weights=weights_1, mdtol=mdtol) var_1 = (s1**2).collapsed(corr_coords, iris.analysis.SUM, weights=weights_1) var_2 = (s2**2).collapsed(corr_coords, iris.analysis.SUM, weights=weights_2) denom = iris.analysis.maths.apply_ufunc(np.sqrt, var_1*var_2, new_unit=covar.units) corr_cube = covar / denom corr_cube.rename("Pearson's r") return corr_cube