def interpolate(cube, sample_points, method=None): """ Extract a sub-cube at the given n-dimensional points. Args: * cube The source Cube. * sample_points A sequence of coordinate (name) - values pairs. Kwargs: * method Request "linear" interpolation (default) or "nearest" neighbour. Only nearest neighbour is available when specifying multi-dimensional coordinates. For example:: sample_points = [('latitude', [45, 45, 45]), ('longitude', [-60, -50, -40])] interpolated_cube = interpolate(cube, sample_points) """ if method not in [None, "linear", "nearest"]: raise ValueError("Unhandled interpolation specified : %s" % method) # Convert any coordinate names to coords points = [] for coord, values in sample_points: if isinstance(coord, six.string_types): coord = cube.coord(coord) points.append((coord, values)) sample_points = points # Do all value sequences have the same number of values? coord, values = sample_points[0] trajectory_size = len(values) for coord, values in sample_points[1:]: if len(values) != trajectory_size: raise ValueError('Lengths of coordinate values are inconsistent.') # Which dimensions are we squishing into the last dimension? squish_my_dims = set() for coord, values in sample_points: dims = cube.coord_dims(coord) for dim in dims: squish_my_dims.add(dim) # Derive the new cube's shape by filtering out all the dimensions we're # about to sample, # and then adding a new dimension to accommodate all the sample points. remaining = [(dim, size) for dim, size in enumerate(cube.shape) if dim not in squish_my_dims] new_data_shape = [size for dim, size in remaining] new_data_shape.append(trajectory_size) # Start with empty data and then fill in the "column" of values for each # trajectory point. new_cube = iris.cube.Cube(np.empty(new_data_shape)) new_cube.metadata = cube.metadata # Derive the mapping from the non-trajectory source dimensions to their # corresponding destination dimensions. remaining_dims = [dim for dim, size in remaining] dimension_remap = {dim: i for i, dim in enumerate(remaining_dims)} # Record a mapping from old coordinate IDs to new coordinates, # for subsequent use in creating updated aux_factories. coord_mapping = {} # Create all the non-squished coords for coord in cube.dim_coords: src_dims = cube.coord_dims(coord) if squish_my_dims.isdisjoint(src_dims): dest_dims = [dimension_remap[dim] for dim in src_dims] new_coord = coord.copy() new_cube.add_dim_coord(new_coord, dest_dims) coord_mapping[id(coord)] = new_coord for coord in cube.aux_coords: src_dims = cube.coord_dims(coord) if squish_my_dims.isdisjoint(src_dims): dest_dims = [dimension_remap[dim] for dim in src_dims] new_coord = coord.copy() new_cube.add_aux_coord(new_coord, dest_dims) coord_mapping[id(coord)] = new_coord # Create all the squished (non derived) coords, not filled in yet. trajectory_dim = len(remaining_dims) for coord in cube.dim_coords + cube.aux_coords: src_dims = cube.coord_dims(coord) if not squish_my_dims.isdisjoint(src_dims): points = np.array([coord.points.flatten()[0]] * trajectory_size) new_coord = iris.coords.AuxCoord(points, standard_name=coord.standard_name, long_name=coord.long_name, units=coord.units, bounds=None, attributes=coord.attributes, coord_system=coord.coord_system) new_cube.add_aux_coord(new_coord, trajectory_dim) coord_mapping[id(coord)] = new_coord for factory in cube.aux_factories: new_cube.add_aux_factory(factory.updated(coord_mapping)) # Are the given coords all 1-dimensional? (can we do linear interp?) for coord, values in sample_points: if coord.ndim > 1: if method == "linear": msg = "Cannot currently perform linear interpolation for " \ "multi-dimensional coordinates." raise iris.exceptions.CoordinateMultiDimError(msg) method = "nearest" break if method in ["linear", None]: for i in range(trajectory_size): point = [(coord, values[i]) for coord, values in sample_points] column = linear_regrid(cube, point) new_cube.data[..., i] = column.data # Fill in the empty squashed (non derived) coords. for column_coord in column.dim_coords + column.aux_coords: src_dims = cube.coord_dims(column_coord) if not squish_my_dims.isdisjoint(src_dims): if len(column_coord.points) != 1: msg = "Expected to find exactly one point. Found {}." raise Exception(msg.format(column_coord.points)) new_cube.coord(column_coord.name()).points[i] = \ column_coord.points[0] elif method == "nearest": # Use a cache with _nearest_neighbour_indices_ndcoords() cache = {} column_indexes = _nearest_neighbour_indices_ndcoords(cube, sample_points, cache=cache) # Construct "fancy" indexes, so we can create the result data array in # a single numpy indexing operation. # ALSO: capture the index range in each dimension, so that we can fetch # only a required (square) sub-region of the source data. fancy_source_indices = [] region_slices = [] n_index_length = len(column_indexes[0]) dims_reduced = [False] * n_index_length for i_ind in range(n_index_length): contents = [column_index[i_ind] for column_index in column_indexes] each_used = [content != slice(None) for content in contents] if np.all(each_used): # This dimension is addressed : use a list of indices. dims_reduced[i_ind] = True # Select the region by min+max indices. start_ind = np.min(contents) stop_ind = 1 + np.max(contents) region_slice = slice(start_ind, stop_ind) # Record point indices with start subtracted from all of them. fancy_index = list(np.array(contents) - start_ind) elif not np.any(each_used): # This dimension is not addressed by the operation. # Use a ":" as the index. fancy_index = slice(None) # No sub-region selection for this dimension. region_slice = slice(None) else: # Should really never happen, if _ndcoords is right. msg = ('Internal error in trajectory interpolation : point ' 'selection indices should all have the same form.') raise ValueError(msg) fancy_source_indices.append(fancy_index) region_slices.append(region_slice) # Fetch the required (square-section) region of the source data. # NOTE: This is not quite as good as only fetching the individual # points used, but it avoids creating a sub-cube for each point, # which is very slow, especially when points are re-used a lot ... source_area_indices = tuple(region_slices) source_data = cube[source_area_indices].data # Transpose source data before indexing it to get the final result. # Because.. the fancy indexing will replace the indexed (horizontal) # dimensions with a new single dimension over trajectory points. # Move those dimensions to the end *first* : this ensures that the new # dimension also appears at the end, which is where we want it. # Make a list of dims with the reduced ones last. dims_reduced = np.array(dims_reduced) dims_order = np.arange(n_index_length) dims_order = np.concatenate( (dims_order[~dims_reduced], dims_order[dims_reduced])) # Rearrange the data dimensions and the fancy indices into that order. source_data = source_data.transpose(dims_order) fancy_source_indices = [ fancy_source_indices[i_dim] for i_dim in dims_order ] # Apply the fancy indexing to get all the result data points. source_data = source_data[fancy_source_indices] # "Fix" problems with missing datapoints producing odd values # when copied from a masked into an unmasked array. # TODO: proper masked data handling. if np.ma.isMaskedArray(source_data): # This is **not** proper mask handling, because we cannot produce a # masked result, but it ensures we use a "filled" version of the # input in this case. source_data = source_data.filled() new_cube.data[:] = source_data # NOTE: we assign to "new_cube.data[:]" and *not* just "new_cube.data", # because the existing code produces a default dtype from 'np.empty' # instead of preserving the input dtype. # TODO: maybe this should be fixed -- i.e. to preserve input dtype ?? # Fill in the empty squashed (non derived) coords. column_coords = [ coord for coord in cube.dim_coords + cube.aux_coords if not squish_my_dims.isdisjoint(cube.coord_dims(coord)) ] new_cube_coords = [ new_cube.coord(column_coord.name()) for column_coord in column_coords ] all_point_indices = np.array(column_indexes) single_point_test_cube = cube[column_indexes[0]] for new_cube_coord, src_coord in zip(new_cube_coords, column_coords): # Check structure of the indexed coord (at one selected point). point_coord = single_point_test_cube.coord(src_coord) if len(point_coord.points) != 1: msg = ('Coord {} at one x-y position has the shape {}, ' 'instead of being a single point. ') raise ValueError(msg.format(src_coord.name(), src_coord.shape)) # Work out which indices apply to the input coord. # NOTE: we know how to index the source cube to get a cube with a # single point for each coord, but this is very inefficient. # So here, we translate cube indexes into *coord* indexes. src_coord_dims = cube.coord_dims(src_coord) fancy_coord_index_arrays = [ list(all_point_indices[:, src_dim]) for src_dim in src_coord_dims ] # Fill the new coord with all the correct points from the old one. new_cube_coord.points = src_coord.points[fancy_coord_index_arrays] # NOTE: the new coords do *not* have bounds. return new_cube
def interpolate(cube, sample_points, method=None): """ Extract a sub-cube at the given n-dimensional points. Args: * cube The source Cube. * sample_points A sequence of coordinate (name) - values pairs. Kwargs: * method Request "linear" interpolation (default) or "nearest" neighbour. Only nearest neighbour is available when specifying multi-dimensional coordinates. For example:: sample_points = [('latitude', [45, 45, 45]), ('longitude', [-60, -50, -40])] interpolated_cube = interpolate(cube, sample_points) """ if method not in [None, "linear", "nearest"]: raise ValueError("Unhandled interpolation specified : %s" % method) # Convert any coordinate names to coords points = [] for coord, values in sample_points: if isinstance(coord, six.string_types): coord = cube.coord(coord) points.append((coord, values)) sample_points = points # Do all value sequences have the same number of values? coord, values = sample_points[0] trajectory_size = len(values) for coord, values in sample_points[1:]: if len(values) != trajectory_size: raise ValueError('Lengths of coordinate values are inconsistent.') # Which dimensions are we squishing into the last dimension? squish_my_dims = set() for coord, values in sample_points: dims = cube.coord_dims(coord) for dim in dims: squish_my_dims.add(dim) # Derive the new cube's shape by filtering out all the dimensions we're # about to sample, # and then adding a new dimension to accommodate all the sample points. remaining = [(dim, size) for dim, size in enumerate(cube.shape) if dim not in squish_my_dims] new_data_shape = [size for dim, size in remaining] new_data_shape.append(trajectory_size) # Start with empty data and then fill in the "column" of values for each # trajectory point. new_cube = iris.cube.Cube(np.empty(new_data_shape)) new_cube.metadata = cube.metadata # Derive the mapping from the non-trajectory source dimensions to their # corresponding destination dimensions. remaining_dims = [dim for dim, size in remaining] dimension_remap = {dim: i for i, dim in enumerate(remaining_dims)} # Record a mapping from old coordinate IDs to new coordinates, # for subsequent use in creating updated aux_factories. coord_mapping = {} # Create all the non-squished coords for coord in cube.dim_coords: src_dims = cube.coord_dims(coord) if squish_my_dims.isdisjoint(src_dims): dest_dims = [dimension_remap[dim] for dim in src_dims] new_coord = coord.copy() new_cube.add_dim_coord(new_coord, dest_dims) coord_mapping[id(coord)] = new_coord for coord in cube.aux_coords: src_dims = cube.coord_dims(coord) if squish_my_dims.isdisjoint(src_dims): dest_dims = [dimension_remap[dim] for dim in src_dims] new_coord = coord.copy() new_cube.add_aux_coord(new_coord, dest_dims) coord_mapping[id(coord)] = new_coord # Create all the squished (non derived) coords, not filled in yet. trajectory_dim = len(remaining_dims) for coord in cube.dim_coords + cube.aux_coords: src_dims = cube.coord_dims(coord) if not squish_my_dims.isdisjoint(src_dims): points = np.array([coord.points.flatten()[0]] * trajectory_size) new_coord = iris.coords.AuxCoord(points, standard_name=coord.standard_name, long_name=coord.long_name, units=coord.units, bounds=None, attributes=coord.attributes, coord_system=coord.coord_system) new_cube.add_aux_coord(new_coord, trajectory_dim) coord_mapping[id(coord)] = new_coord for factory in cube.aux_factories: new_cube.add_aux_factory(factory.updated(coord_mapping)) # Are the given coords all 1-dimensional? (can we do linear interp?) for coord, values in sample_points: if coord.ndim > 1: if method == "linear": msg = "Cannot currently perform linear interpolation for " \ "multi-dimensional coordinates." raise iris.exceptions.CoordinateMultiDimError(msg) method = "nearest" break if method in ["linear", None]: for i in range(trajectory_size): point = [(coord, values[i]) for coord, values in sample_points] column = linear_regrid(cube, point) new_cube.data[..., i] = column.data # Fill in the empty squashed (non derived) coords. for column_coord in column.dim_coords + column.aux_coords: src_dims = cube.coord_dims(column_coord) if not squish_my_dims.isdisjoint(src_dims): if len(column_coord.points) != 1: msg = "Expected to find exactly one point. Found {}." raise Exception(msg.format(column_coord.points)) new_cube.coord(column_coord.name()).points[i] = \ column_coord.points[0] elif method == "nearest": # Use a cache with _nearest_neighbour_indices_ndcoords() cache = {} column_indexes = _nearest_neighbour_indices_ndcoords( cube, sample_points, cache=cache) # Construct "fancy" indexes, so we can create the result data array in # a single numpy indexing operation. # ALSO: capture the index range in each dimension, so that we can fetch # only a required (square) sub-region of the source data. fancy_source_indices = [] region_slices = [] n_index_length = len(column_indexes[0]) dims_reduced = [False] * n_index_length for i_ind in range(n_index_length): contents = [column_index[i_ind] for column_index in column_indexes] each_used = [content != slice(None) for content in contents] if np.all(each_used): # This dimension is addressed : use a list of indices. dims_reduced[i_ind] = True # Select the region by min+max indices. start_ind = np.min(contents) stop_ind = 1 + np.max(contents) region_slice = slice(start_ind, stop_ind) # Record point indices with start subtracted from all of them. fancy_index = list(np.array(contents) - start_ind) elif not np.any(each_used): # This dimension is not addressed by the operation. # Use a ":" as the index. fancy_index = slice(None) # No sub-region selection for this dimension. region_slice = slice(None) else: # Should really never happen, if _ndcoords is right. msg = ('Internal error in trajectory interpolation : point ' 'selection indices should all have the same form.') raise ValueError(msg) fancy_source_indices.append(fancy_index) region_slices.append(region_slice) # Fetch the required (square-section) region of the source data. # NOTE: This is not quite as good as only fetching the individual # points used, but it avoids creating a sub-cube for each point, # which is very slow, especially when points are re-used a lot ... source_area_indices = tuple(region_slices) source_data = cube[source_area_indices].data # Transpose source data before indexing it to get the final result. # Because.. the fancy indexing will replace the indexed (horizontal) # dimensions with a new single dimension over trajectory points. # Move those dimensions to the end *first* : this ensures that the new # dimension also appears at the end, which is where we want it. # Make a list of dims with the reduced ones last. dims_reduced = np.array(dims_reduced) dims_order = np.arange(n_index_length) dims_order = np.concatenate((dims_order[~dims_reduced], dims_order[dims_reduced])) # Rearrange the data dimensions and the fancy indices into that order. source_data = source_data.transpose(dims_order) fancy_source_indices = [fancy_source_indices[i_dim] for i_dim in dims_order] # Apply the fancy indexing to get all the result data points. source_data = source_data[fancy_source_indices] # "Fix" problems with missing datapoints producing odd values # when copied from a masked into an unmasked array. # TODO: proper masked data handling. if np.ma.isMaskedArray(source_data): # This is **not** proper mask handling, because we cannot produce a # masked result, but it ensures we use a "filled" version of the # input in this case. source_data = source_data.filled() new_cube.data[:] = source_data # NOTE: we assign to "new_cube.data[:]" and *not* just "new_cube.data", # because the existing code produces a default dtype from 'np.empty' # instead of preserving the input dtype. # TODO: maybe this should be fixed -- i.e. to preserve input dtype ?? # Fill in the empty squashed (non derived) coords. column_coords = [coord for coord in cube.dim_coords + cube.aux_coords if not squish_my_dims.isdisjoint( cube.coord_dims(coord))] new_cube_coords = [new_cube.coord(column_coord.name()) for column_coord in column_coords] all_point_indices = np.array(column_indexes) single_point_test_cube = cube[column_indexes[0]] for new_cube_coord, src_coord in zip(new_cube_coords, column_coords): # Check structure of the indexed coord (at one selected point). point_coord = single_point_test_cube.coord(src_coord) if len(point_coord.points) != 1: msg = ('Coord {} at one x-y position has the shape {}, ' 'instead of being a single point. ') raise ValueError(msg.format(src_coord.name(), src_coord.shape)) # Work out which indices apply to the input coord. # NOTE: we know how to index the source cube to get a cube with a # single point for each coord, but this is very inefficient. # So here, we translate cube indexes into *coord* indexes. src_coord_dims = cube.coord_dims(src_coord) fancy_coord_index_arrays = [list(all_point_indices[:, src_dim]) for src_dim in src_coord_dims] # Fill the new coord with all the correct points from the old one. new_cube_coord.points = src_coord.points[fancy_coord_index_arrays] # NOTE: the new coords do *not* have bounds. return new_cube
def interpolate(cube, sample_points, method=None): """ Extract a sub-cube at the given n-dimensional points. Args: * cube The source Cube. * sample_points A sequence of coordinate (name) - values pairs. Kwargs: * method Request "linear" interpolation (default) or "nearest" neighbour. Only nearest neighbour is available when specifying multi-dimensional coordinates. For example:: sample_points = [('latitude', [45, 45, 45]), ('longitude', [-60, -50, -40])] interpolated_cube = interpolate(cube, sample_points) """ if method not in [None, "linear", "nearest"]: raise ValueError("Unhandled interpolation specified : %s" % method) # Convert any coordinate names to coords points = [] for coord, values in sample_points: if isinstance(coord, six.string_types): coord = cube.coord(coord) points.append((coord, values)) sample_points = points # Do all value sequences have the same number of values? coord, values = sample_points[0] trajectory_size = len(values) for coord, values in sample_points[1:]: if len(values) != trajectory_size: raise ValueError('Lengths of coordinate values are inconsistent.') # Which dimensions are we squishing into the last dimension? squish_my_dims = set() for coord, values in sample_points: dims = cube.coord_dims(coord) for dim in dims: squish_my_dims.add(dim) # Derive the new cube's shape by filtering out all the dimensions we're about to sample, # and then adding a new dimension to accommodate all the sample points. remaining = [(dim, size) for dim, size in enumerate(cube.shape) if dim not in squish_my_dims] new_data_shape = [size for dim, size in remaining] new_data_shape.append(trajectory_size) # Start with empty data and then fill in the "column" of values for each trajectory point. new_cube = iris.cube.Cube(np.empty(new_data_shape)) new_cube.metadata = cube.metadata # Derive the mapping from the non-trajectory source dimensions to their # corresponding destination dimensions. remaining_dims = [dim for dim, size in remaining] dimension_remap = {dim: i for i, dim in enumerate(remaining_dims)} # Record a mapping from old coordinate IDs to new coordinates, # for subsequent use in creating updated aux_factories. coord_mapping = {} # Create all the non-squished coords for coord in cube.dim_coords: src_dims = cube.coord_dims(coord) if squish_my_dims.isdisjoint(src_dims): dest_dims = [dimension_remap[dim] for dim in src_dims] new_coord = coord.copy() new_cube.add_dim_coord(new_coord, dest_dims) coord_mapping[id(coord)] = new_coord for coord in cube.aux_coords: src_dims = cube.coord_dims(coord) if squish_my_dims.isdisjoint(src_dims): dest_dims = [dimension_remap[dim] for dim in src_dims] new_coord = coord.copy() new_cube.add_aux_coord(new_coord, dest_dims) coord_mapping[id(coord)] = new_coord # Create all the squished (non derived) coords, not filled in yet. trajectory_dim = len(remaining_dims) for coord in cube.dim_coords + cube.aux_coords: src_dims = cube.coord_dims(coord) if not squish_my_dims.isdisjoint(src_dims): points = np.array([coord.points.flatten()[0]] * trajectory_size) new_coord = iris.coords.AuxCoord(points, standard_name=coord.standard_name, long_name=coord.long_name, units=coord.units, bounds=None, attributes=coord.attributes, coord_system=coord.coord_system) new_cube.add_aux_coord(new_coord, trajectory_dim) coord_mapping[id(coord)] = new_coord for factory in cube.aux_factories: new_cube.add_aux_factory(factory.updated(coord_mapping)) # Are the given coords all 1-dimensional? (can we do linear interp?) for coord, values in sample_points: if coord.ndim > 1: if method == "linear": raise iris.exceptions.CoordinateMultiDimError("Cannot currently perform linear interpolation for multi-dimensional coordinates.") method = "nearest" break # Use a cache with _nearest_neighbour_indices_ndcoords() cache = {} for i in range(trajectory_size): point = [(coord, values[i]) for coord, values in sample_points] if method in ["linear", None]: column = linear_regrid(cube, point) new_cube.data[..., i] = column.data elif method == "nearest": column_index = _nearest_neighbour_indices_ndcoords(cube, point, cache=cache) column = cube[column_index] new_cube.data[..., i] = column.data # Fill in the empty squashed (non derived) coords. for column_coord in column.dim_coords + column.aux_coords: src_dims = cube.coord_dims(column_coord) if not squish_my_dims.isdisjoint(src_dims): if len(column_coord.points) != 1: raise Exception("Expected to find exactly one point. Found %d" % len(column_coord.points)) new_cube.coord(column_coord.name()).points[i] = column_coord.points[0] return new_cube