This document is describing generally the present state of Quantity, to present to the gt4py team. It is meant to communicate some general use cases of Quantity and the sort of metadata which has been useful for it to provide, rather than to describe how to use a particular implementation.
API
Most of the bullets below contain attribute descriptions using Python-style type hinting.
Quantity
Attributes required on initialization:
data(ndarray-like): the underlying numpy or cupy array(or storage, presently)
dims(Iterable[str]): names of each dimension
units(str)
origin(Iterable[int]): the start of the compute domain
extent(Iterable[int]): the shape of the compute domain
data_array(xr.DataArray): xarray-wrapped reference to compute domain data
values(np.ndarray): a read-only numpy array referencing the data
np: numpy-like module to interact with the underlying data
view(BoundedArrayView): provides indexing into the underlying data offset by the origin with the compute bounds as default, e.g. quantity.view[:, :, :] provides the first 3 compute points, quantity.view[-1, -1, :] provides the southwest(bottom left) corner halo point.
This attribute is being re-worked into a set of attributes providing different indexing offset behaviors:
interior for the compute domain
north, east, south, west for edges
northeast, northwest, southeast, southwest for corners
less important attributes(for our discussion):
metadata(QuantityMetadata): structure containing info required to initialize a matching Quantity object for MPI receives
attrs(Mapping[str, str]): compatibility attribute for DataArray, contains units information
Usage examples
Zero the compute domain of a quantity:
quantity.view[:] = 0.
Rotation with arbitrary dimension order(used in halo updates):
API
Usage examples
quantity.view[:] = 0.def rotate_scalar_data(data, dims, numpy, n_clockwise_rotations): n_clockwise_rotations = n_clockwise_rotations % 4 if n_clockwise_rotations == 0: pass elif n_clockwise_rotations in (1, 3): x_dim, y_dim = None, None for i, dim in enumerate(dims): if dim in constants.X_DIMS: x_dim = i elif dim in constants.Y_DIMS: y_dim = i if (x_dim is not None) and (y_dim is not None): if n_clockwise_rotations == 1: data = numpy.rot90(data, axes=(y_dim, x_dim)) elif n_clockwise_rotations == 3: data = numpy.rot90(data, axes=(x_dim, y_dim)) elif (x_dim is None) != (y_dim is None): # requires converting y dims to x dims and vice-versa raise NotImplementedError( "cannot yet rotate values which don't have both x and y dims" ) elif n_clockwise_rotations == 2: slice_list = [] for dim in dims: if dim in constants.HORIZONTAL_DIMS: slice_list.append(slice(None, None, -1)) else: slice_list.append(slice(None, None)) data = data[tuple(slice_list)] return data def get_origin(self, dims: Tuple[str, ...]) -> Tuple[int, ...]: return_list = [ self.n_halo if dim in constants.HORIZONTAL_DIMS else 0 for dim in dims ] return tuple(return_list)
def get_extent(self, dims: Tuple[str, ...]) -> Tuple[int, ...]: extents = self.extra_dim_lengths.copy() extents.update( { constants.X_DIM: self.nx, constants.X_INTERFACE_DIM: self.nx + 1, constants.Y_DIM: self.ny, constants.Y_INTERFACE_DIM: self.ny + 1, constants.Z_DIM: self.nz, constants.Z_INTERFACE_DIM: self.nz + 1, } return tuple(extents[dim] for dim in dims)
def get_shape(self, dims: Tuple[str, ...]) -> Tuple[int, ...]: shape_dict = self.extra_dim_lengths.copy() # must pad non-interface variables to have the same shape as interface variables shape_dict.update( { constants.X_DIM: self.nx + 1 + 2 * self.n_halo, constants.X_INTERFACE_DIM: self.nx + 1 + 2 * self.n_halo, constants.Y_DIM: self.ny + 1 + 2 * self.n_halo, constants.Y_INTERFACE_DIM: self.ny + 1 + 2 * self.n_halo, constants.Z_DIM: self.nz + 1, constants.Z_INTERFACE_DIM: self.nz + 1, } ) return tuple(shape_dict[dim] for dim in dims)