utils

GPU_AVAILABLE = True module-attribute

T = TypeVar('T') module-attribute

DEFAULT_GRID_NML_GROUPS = ['fv_core_nml'] module-attribute

MetaEnumStr

Bases: EnumMeta

list_by_dims(dims, horizontal_list, non_horizontal_value)

Take in a list of dimensions, a (y, x) set of values, and a value for any non-horizontal dimensions. Return a list of length len(dims) with the value for each dimension.

is_contiguous(array)

is_c_contiguous(array)

ensure_contiguous(maybe_array)

safe_assign_array(to_array, from_array)

Failproof assignment for array on different devices.

The memory will be downloaded/uploaded from GPU if need be.

Parameters:

Name	Type	Description	Default
to_array	ndarray	destination ndarray	required
from_array	ndarray	source ndarray	required

device_synchronize()

Synchronize all memory communication

safe_mpi_allocate(allocator, shape, dtype)

Make sure the allocation use an allocator that works with MPI

For G2G transfer, MPICH requires the allocation to not be done with managed memory. Since we can't know what state cupy is in with switch for the default pooled allocator.

If allocator comes from cupy, it must be cupy.empty or cupy.zeros. We raise a RuntimeError if a cupy array is allocated outside of the safe code path.

Though the allocation might be safe, the MPI crash that result from a managed memory allocation is non trivial and should be tightly controlled.

flatten_nml_to_dict(nml)

Returns a flattened dict version of a f90nml.namelist.Namelist

Parameters:

Name	Type	Description	Default
nml	Namelist	f90nml.Namelist	required

load_f90nml(namelist_path)

Loads a Fortran namelist given its path and return a f90nml.Namelist

Parameters:

Name	Type	Description	Default
namelist_path	Path	Path to the Fortran namelist file	required

load_f90nml_as_dict(namelist_path, flatten=True, target_groups=None)

Loads a Fortran namelist given its path and returns a dict representation. If target_groups are specified, then the dict is created using only those groups.

Parameters:

Name	Type	Description	Default
namelist_path	Path	Path to the Fortran namelist file	required
flatten	bool	If True, flattens the loaded namelist (without groups) before returning it. (Default: True) Otherwise, it returns the f90nml.Namelist dict representation.	True
target_groups	list[str] | None	If 'None' is specified, then all groups are considered. (Default: None) Otherwise, only parameters from the specified groups are considered.	None

f90nml_as_dict(nml, flatten=True, target_groups=None)

Uses a f90nml.Namelist and returns a dict representation. If target_groups are specified, then the dict is created using only those groups. The return dicts can be flattened further to remove the group information or keep the group information.

Parameters:

Name	Type	Description	Default
nml	Namelist	f90nml.Namelist	required
flatten	bool	If True, flattens the loaded namelist (without groups) before returning it. (Default: True) Otherwise, it returns the f90nml.Namelist dict representation.	True
target_groups	list[str] | None	If 'None' is specified, then all groups are considered. (Default: None) Otherwise, only parameters from the specified groups are considered.	None

grid_params_from_f90nml(nml)

Uses a f90nml.Namelist and returns a dict representation of parameters useful for grid generation. The return dict will be flattened with key-value pairs from the nml's DEFAULT_GRID_NML_GROUPS.

Parameters:

Name	Type	Description	Default
nml	Namelist	f90nml.Namelist	required