gt4py_utils

managed_memory = True module-attribute

origin = (N_HALO_DEFAULT, N_HALO_DEFAULT, 0) module-attribute

tracer_variables = ['qvapor', 'qliquid', 'qrain', 'qice', 'qsnow', 'qgraupel', 'qo3mr', 'qsgs_tke', 'qcld'] module-attribute

mark_untested(msg='This is not tested')

make_storage_data(data, shape=None, origin=origin, *, backend, dtype=Float, mask=None, start=(0, 0, 0), dummy=None, axis=2, max_dim=3, read_only=True)

Create a new gt4py storage from the given data.

Parameters:

Name	Type	Description	Default
data	NDArray	Data array for new storage	required
shape	tuple[int, ...] | None	Shape of the new storage. Number of indices should be equal to number of unmasked axes	None
origin	tuple[int, ...]	Default origin for gt4py stencil calls	origin
dtype	DTypeLike	Data type	Float
mask	tuple[bool, ...] | None	Tuple indicating the axes used when initializing the storage. True indicates a masked axis, False is a used axis.	None
start	tuple[int, ...]	Starting points for slices in data copies	(0, 0, 0)
dummy	tuple[int, ...] | None	Dummy axes	None
axis	int	Axis for 2D to 3D arrays	2
backend	Backend	current backend in use	required

Returns:

Type	Description
NDArray	Field[..., dtype]: New storage

Examples:

1) ptop = utils.make_storage_data(top_p, q4_1.shape) 2) ws3 = utils.make_storage_data(ws3[:, :, -1], shape, origin=(0, 0, 0)) 3) data_dict[names[i]] = make_storage_data( data[:, :, :, i], shape, origin=origin, start=start, dummy=dummy, axis=axis, )

make_storage_from_shape(shape, origin=origin, *, backend, dtype=Float, mask=None)

Create a new gt4py storage of a given shape filled with zeros.

Parameters:

Name	Type	Description	Default
shape	tuple[int, ...]	Shape of the new storage	required
origin	tuple[int, ...]	Default origin for gt4py stencil calls	origin
dtype	DTypeLike	Data type	Float
mask	tuple[bool, ...] | None	Tuple indicating the axes used when initializing the storage	None
backend	Backend	gt4py backend to use when making the storage	required

Returns:

Type	Description
NDArray	Field[..., dtype]: New storage

Examples:

1) utmp = utils.make_storage_from_shape(ua.shape) 2) qx = utils.make_storage_from_shape( qin.shape, origin=(grid().is_, grid().jsd, kstart) ) 3) q_out = utils.make_storage_from_shape(q_in.shape, origin,)

make_storage_dict(data, shape=None, origin=origin, start=(0, 0, 0), dummy=None, names=None, axis=2, *, backend, dtype=Float)

storage_dict(st_dict, names, shape, origin, *, backend)

get_kstarts(column_info, npz)

k_split_run(func, data, k_indices, splitvars_values)

asarray(array, to_type=np.ndarray, dtype=None, order=None)

zeros(shape, dtype=Float, *, backend)

sum(array, axis=None, dtype=Float, out=None, keepdims=False)

repeat(array, repeats, axis=None)

index(array, key)

moveaxis(array, source, destination)

tile(array, reps)

squeeze(array, axis=None)

reshape(array, new_shape)

unique(array, return_index=False, return_inverse=False, return_counts=False, axis=None)

stack(tup, axis=0, out=None)

device_sync(backend)

split_cartesian_into_storages(var)

Provided a storage of dims [I_DIM, J_DIM, CARTESIAN_DIM] or [I_INTERFACE_DIM, J_INTERFACE_DIM, CARTESIAN_DIM] Split it into separate 2D storages for each cartesian dimension, and return these in a list.