def transpose(state, dims, npz, npx, npy):
return_state = {}
for name, quantity in state.items():
if name == "time":
return_state[name] = quantity
else:
if len(quantity.storage.shape) == 2:
data_3d = numpy.ascontiguousarray(
numpy.broadcast_to(
quantity.data[:, :, None],
(quantity.data.shape[0], quantity.data.shape[1],
npz + 1),
))
quantity_3d = Quantity.from_data_array(
xr.DataArray(
data_3d,
attrs=quantity.attrs,
dims=[
quantity.dims[0], quantity.dims[1], Z_INTERFACE_DIM
],
),
origin=(quantity.origin[0], quantity.origin[1], 0),
extent=(quantity.extent[0], quantity.extent[1], npz + 1),
)
quantity_3d.metadata.gt4py_backend = "numpy"
return_state[name] = quantity_3d.transpose(dims)
elif len(quantity.storage.shape) == 1:
data_3d = numpy.tile(quantity.data, (npx + 6, npy + 6, 1))
quantity_3d = Quantity.from_data_array(
xr.DataArray(
data_3d,
attrs=quantity.attrs,
dims=[
X_INTERFACE_DIM, Y_INTERFACE_DIM, quantity.dims[0]
],
),
origin=(0, 0, quantity.origin[0]),
extent=(npx, npy, quantity.extent[0]),
)
quantity_3d.metadata.gt4py_backend = "numpy"
return_state[name] = quantity_3d.transpose(dims)
else:
return_state[name] = quantity.transpose(dims)
return return_state
def _convert_to_quantities(state, metadata):
quantities = {}
for key, data in state.items():
data = np.squeeze(data.astype(np.float32))
data_t = data.T
dims = DIMS_MAP[data.ndim]
attrs = _get_attrs(key, metadata)
units = attrs.pop("units", "unknown")
quantities[key] = Quantity(data_t, dims, units)
# Access to private member could break TODO: Quantity kwarg for attrs?
quantities[key]._attrs.update(attrs)
return quantities
def convert_3d_to_1d(state, field_names):
return_state = state
for field in field_names:
quantity = state[field]
# Assuming we've already transposed from xyz to zyx
data_1d = quantity.data[
:, 0, 0
] # take the first column since they should be the same
quantity_1d = Quantity.from_data_array(
xr.DataArray(data_1d, attrs=quantity.attrs, dims=[quantity.dims[0]]),
origin=[quantity.origin[0]],
extent=[quantity.extent[0]],
)
return_state[field] = quantity_1d
return return_state
def convert_3d_to_2d(state, field_levels):
return_state = state
for field in field_levels.keys():
quantity = state[field]
# Assuming we've already transposed from xyz to zyx
data_2d = quantity.data[field_levels[
field], :, :] # take the bottom level since they should all be the same
quantity_2d = Quantity.from_data_array(
xr.DataArray(data_2d,
attrs=quantity.attrs,
dims=[quantity.dims[1], quantity.dims[2]]),
origin=(quantity.origin[1], quantity.origin[2]),
extent=(quantity.extent[1], quantity.extent[2]),
)
return_state[field] = quantity_2d
return return_state
fv3core._config.set_grid(grid)
# startup
wrapper.initialize()
# add things to State
origin = (0, 3, 3)
extent = (spec.namelist.npz, spec.namelist.npy - 1, spec.namelist.npx - 1)
arr = np.zeros(
(spec.namelist.npz + 1, spec.namelist.npy + 6, spec.namelist.npx + 6))
turbulent_kinetic_energy = Quantity.from_data_array(
xr.DataArray(
arr,
attrs={
"fortran_name": "qsgs_tke",
"units": "m**2/s**2"
},
dims=["z", "y", "x"],
),
origin=origin,
extent=extent,
)
u_tendency = Quantity.from_data_array(
xr.DataArray(
arr.reshape((spec.namelist.npx + 6, spec.namelist.npy + 6,
spec.namelist.npz + 1)),
attrs={
"fortran_name": "u_dt",
"units": "m/s**2"
},
dims=["x", "y", "z"],
),