API¶

exception fv3gfs.util.InvalidQuantityError¶: Bases: Exception

exception fv3gfs.util.OutOfBoundsError¶: Bases: ValueError

fv3gfs.util.datetime64_to_datetime(dt64: numpy.datetime64) → datetime.datetime¶

fv3gfs.util.read_state(filename: str) → dict¶

Read a model state from a NetCDF file.

Parameters:	filename – local or remote location of the NetCDF file
Returns:	a model state dictionary
Return type:	state

fv3gfs.util.write_state(state: dict, filename: str) → None¶

Write a model state to a NetCDF file.

Parameters:	state – a model state dictionary filename – local or remote location to write the NetCDF file

fv3gfs.util.get_nudging_tendencies(state, reference_state, nudging_timescales: Mapping[str, datetime.timedelta])¶

Return the nudging tendency of the given state towards the reference state according to the provided nudging timescales.

Parameters:

state (dict) – A state dictionary.
reference_state (dict) – A reference state dictionary.
nudging_timescales (dict) – A dictionary whose keys are standard names and values are timedelta objects indicating the relaxation timescale for that variable.

Returns:

A dictionary whose keys are standard names: and values are Quantity objects indicating the nudging tendency of that standard name.

Return type:

nudging_tendencies (dict)

fv3gfs.util.apply_nudging(state, reference_state, nudging_timescales: Mapping[str, datetime.timedelta], timestep: datetime.timedelta)¶

Nudge the given state towards the reference state according to the provided nudging timescales.

Nudging is applied to the state in-place.

Parameters:

state (dict) – A state dictionary.
reference_state (dict) – A reference state dictionary.
nudging_timescales (dict) – A dictionary whose keys are standard names and values are timedelta objects indicating the relaxation timescale for that variable.
timestep (timedelta) – length of the timestep

Returns:

A dictionary whose keys are standard names: and values are Quantity objects indicating the nudging tendency of that standard name.

Return type:

nudging_tendencies (dict)

fv3gfs.util.open_restart(dirname: str, communicator: fv3gfs.util.communicator.CubedSphereCommunicator, label: str = '', only_names: Iterable[str] = None, to_state: dict = None, tracer_properties: Mapping[str, Mapping[str, Union[str, Iterable[str]]]] = None)¶

Load restart files output by the Fortran model into a state dictionary.

Parameters:	dirname – location of restart files, can be local or remote communicator – object for communication over the cubed sphere label – prepended string on the restart files to load only_names (optional) – list of standard names to load to_state (optional) – if given, assign loaded data into pre-allocated quantities in this state dictionary
Returns:	model state dictionary
Return type:	state

class fv3gfs.util.ZarrMonitor(store: Union[str, collections.abc.MutableMapping], partitioner: fv3gfs.util.partitioner.CubedSpherePartitioner, mode: str = 'w', mpi_comm=<fv3gfs.util.zarr_monitor.DummyComm object>)¶

Bases: object

sympl.Monitor-style object for storing model state dictionaries in a Zarr store.

store(state: dict) → None¶

Append the model state dictionary to the zarr store.

Requires the state contain the same quantities with the same metadata as the first time this is called. Quantities are stored with dimensions [time, rank] followed by the dimensions included in any one state snapshot. The one exception is “time” which is stored with dimensions [time].

class fv3gfs.util.CubedSpherePartitioner(tile: fv3gfs.util.partitioner.TilePartitioner)¶

Bases: fv3gfs.util.partitioner.Partitioner

boundary(boundary_type: int, rank: int) → Optional[fv3gfs.util.boundary.SimpleBoundary]¶

Returns a boundary of the requested type for a given rank, or None.

On tile corners, the boundary across that corner does not exist.

Parameters:	boundary_type – the type of boundary rank – the processor rank
Returns:	boundary

classmethod from_namelist(namelist)¶

Initialize a CubedSpherePartitioner from a Fortran namelist.

Parameters:	namelist (dict) – the Fortran namelist

global_extent(rank_metadata: fv3gfs.util.quantity.QuantityMetadata) → Tuple[int, ...]¶

Return the shape of a full cube representation for the given dimensions.

Parameters:	metadata – quantity metadata
Returns:	shape of full cube representation
Return type:	extent

layout¶

subtile_extent(cube_metadata: fv3gfs.util.quantity.QuantityMetadata) → Tuple[int, ...]¶: Return the shape of a single rank representation for the given dimensions.

subtile_slice(rank: int, global_dims: Sequence[str], global_extent: Sequence[int], overlap: bool = False) → Tuple[Union[int, slice], ...]¶

Return the subtile slice of a given rank on an array.

Global refers to the domain being partitioned. For example, for a partitioning of a tile, the tile would be the “global” domain.

Parameters:

rank – the rank of the process
global_dims – dimensions of the global quantity being partitioned
global_extent – extent of the global quantity being partitioned
overlap (optional) – if True, for interface variables include the part of the array shared by adjacent ranks in both ranks. If False, ensure only one of those ranks (the greater rank) is assigned the overlapping section. Default is False.

Returns:

the tuple slice of the global compute domain corresponding: to the subtile compute domain

Return type:

subtile_slice

tile_index(rank: int) → int¶: Returns the tile index of a given rank

tile_root_rank(rank: int) → int¶: Returns the lowest rank on the same tile as a given rank.

total_ranks¶: the number of ranks on the cubed sphere

class fv3gfs.util.TilePartitioner(layout: Tuple[int, int])¶

Bases: fv3gfs.util.partitioner.Partitioner

boundary(boundary_type: int, rank: int) → Optional[fv3gfs.util.boundary.SimpleBoundary]¶

Returns a boundary of the requested type for a given rank.

Target ranks will be on the same tile as the given rank, wrapping around as in a doubly-periodic boundary condition.

Parameters:	boundary_type – the type of boundary rank – the processor rank
Returns:	boundary

fliplr_rank(rank: int) → int¶

classmethod from_namelist(namelist)¶

Initialize a TilePartitioner from a Fortran namelist.

Parameters:	namelist (dict) – the Fortran namelist

global_extent(rank_metadata: fv3gfs.util.quantity.QuantityMetadata) → Tuple[int, ...]¶

Return the shape of a full tile representation for the given dimensions.

Parameters:	metadata – quantity metadata
Returns:	shape of full tile representation
Return type:	extent

on_tile_bottom(rank: int) → bool¶

on_tile_left(rank: int) → bool¶

on_tile_right(rank: int) → bool¶

on_tile_top(rank: int) → bool¶

rotate_rank(rank: int, n_clockwise_rotations: int) → int¶

subtile_extent(global_metadata: fv3gfs.util.quantity.QuantityMetadata) → Tuple[int, ...]¶: Return the shape of a single rank representation for the given dimensions.

subtile_index(rank: int) → Tuple[int, int]¶: Return the (y, x) subtile position of a given rank as an integer number of subtiles.

subtile_slice(rank: int, global_dims: Sequence[str], global_extent: Sequence[int], overlap: bool = False) → Tuple[slice, ...]¶

Return the subtile slice of a given rank on an array.

Global refers to the domain being partitioned. For example, for a partitioning of a tile, the tile would be the “global” domain.

Parameters:

rank – the rank of the process
global_dims – dimensions of the global quantity being partitioned
global_extent – extent of the global quantity being partitioned
overlap (optional) – if True, for interface variables include the part of the array shared by adjacent ranks in both ranks. If False, ensure only one of those ranks (the greater rank) is assigned the overlapping section. Default is False.

Returns:

the slice of the global compute domain corresponding: to the subtile compute domain

Return type:

subtile_slice

total_ranks¶

fv3gfs.util.get_tile_index(rank: int, total_ranks: int) → int¶: Returns the zero-indexed tile number, given a rank and total number of ranks.

fv3gfs.util.get_tile_number(tile_rank: int, total_ranks: int) → int¶

Deprecated: use get_tile_index.

Returns the tile number for a given rank and total number of ranks.

class fv3gfs.util.Timer¶

Bases: object

Class to accumulate timings for named operations.

clock(name: str)¶

Context manager to produce timings of operations.

Parameters:	name – the name of the operation being timed

Example

The context manager times operations that happen within its context. The following would time a time.sleep operation:

>>> import time
>>> from fv3gfs.util import Timer
>>> timer = Timer()
>>> with timer.clock("sleep"):
...     time.sleep(1)
...
>>> timer.times
{'sleep': 1.0032463260000029}

disable()¶: Disable the Timer.

enable()¶: Enable the Timer.

enabled¶: Indicates whether the timer is currently enabled.

hits¶: accumulated hit counts for each operation name

reset()¶: Remove all accumulated timings.

start(name: str)¶: Start timing a given named operation.

stop(name: str)¶: Stop timing a given named operation, add the time elapsed to accumulated timing and increase the hit count.

times¶: accumulated timings for each operation name

class fv3gfs.util.Quantity(data, dims: Sequence[str], units: str, origin: Sequence[int] = None, extent: Sequence[int] = None, gt4py_backend: Optional[str] = None)¶

Bases: object

Data container for physical quantities.

attrs¶

data¶: the underlying array of data

data_array¶

dims¶: names of each dimension

extent¶: the shape of the computational domain

classmethod from_data_array(data_array: xarray.core.dataarray.DataArray, origin: Sequence[int] = None, extent: Sequence[int] = None, gt4py_backend: Optional[str] = None) → fv3gfs.util.quantity.Quantity¶

Initialize a Quantity from an xarray.DataArray.

Parameters:	data_array – origin – first point in data within the computational domain extent – number of points along each axis within the computational domain gt4py_backend – backend to use for gt4py storages, if not given this will be derived from a Storage if given as the data argument, otherwise the storage attribute is disabled and will raise an exception

gt4py_backend¶

metadata¶

np¶

origin¶: the start of the computational domain

sel(**kwargs) → numpy.ndarray¶

Convenience method to perform indexing on view using dimension names without knowing dimension order.

Parameters:	**kwargs – slice/index to retrieve for a given dimension name
Returns:	an ndarray-like selection of the given indices on self.view
Return type:	view_selection

storage¶

A gt4py storage representing the data in this Quantity.

Will raise TypeError if the gt4py backend was not specified when initializing this object, either by providing a Storage for data or explicitly specifying a backend.

transpose(target_dims: Sequence[Union[str, Iterable[str]]]) → fv3gfs.util.quantity.Quantity¶

Change the dimension order of this Quantity.

If you know you are working with cell-centered variables, you can do:

>>> from fv3gfs.util import X_DIM, Y_DIM, Z_DIM
>>> transposed_quantity = quantity.transpose([X_DIM, Y_DIM, Z_DIM])

To support re-ordering without checking whether quantities are on cell centers or interfaces, the API supports giving a list of dimension names for dimensions. For example, to re-order to X-Y-Z dimensions regardless of the grid the variable is on, one could do:

>>> from fv3gfs.util import X_DIMS, Y_DIMS, Z_DIMS
>>> transposed_quantity = quantity.transpose([X_DIMS, Y_DIMS, Z_DIMS])

Parameters:	target_dims – a list of output dimensions. Instead of a single dimension name, an iterable of dimensions can be used instead for any entries. For example, you may want to use fv3gfs.util.X_DIMS to place an x-dimension without knowing whether it is on cell centers or interfaces.
Returns:	Quantity with the requested output dimension order
Return type:	transposed
Raises:	`ValueError` – if any of the target dimensions do not exist on this Quantity, or if this Quantity contains multiple values from an iterable entry

units¶: units of the quantity

values¶

view¶: a view into the computational domain of the underlying data

class fv3gfs.util.QuantityMetadata(origin: Tuple[int, ...], extent: Tuple[int, ...], dims: Tuple[str, ...], units: str, data_type: type, dtype: type, gt4py_backend: Union[str, NoneType] = None)¶

Bases: object

data_type = None¶: ndarray-like type used to store the data

dim_lengths¶: mapping of dimension names to their lengths

dims = None¶: names of each dimension

dtype = None¶: dtype of the data in the ndarray-like object

extent = None¶: the shape of the computational domain

gt4py_backend = None¶: backend to use for gt4py storages

np¶: numpy-like module used to interact with the data

origin = None¶: the start of the computational domain

units = None¶: units of the quantity

fv3gfs.util.ensure_equal_units(units1: str, units2: str) → None¶

fv3gfs.util.units_are_equal(units1: str, units2: str) → bool¶

exception fv3gfs.util.UnitsError¶: Bases: Exception

class fv3gfs.util.TileCommunicator(comm, partitioner: fv3gfs.util.partitioner.TilePartitioner)¶

Bases: fv3gfs.util.communicator.Communicator

Performs communications within a single tile or region of a tile

class fv3gfs.util.CubedSphereCommunicator(comm, partitioner: fv3gfs.util.partitioner.CubedSpherePartitioner)¶

Bases: fv3gfs.util.communicator.Communicator

Performs communications within a cubed sphere

boundaries¶: boundaries of this tile with neighboring tiles

finish_halo_update(quantity: fv3gfs.util.quantity.Quantity, n_points: int)¶: Deprecated, do not use.

finish_vector_halo_update(x_quantity: fv3gfs.util.quantity.Quantity, y_quantity: fv3gfs.util.quantity.Quantity, n_points: int)¶: Deprecated, do not use.

halo_update(quantity: fv3gfs.util.quantity.Quantity, n_points: int)¶

Perform a halo update on a quantity.

Parameters:	quantity – the quantity to be updated n_points – how many halo points to update, starting from the interior

start_halo_update(quantity: fv3gfs.util.quantity.Quantity, n_points: int) → fv3gfs.util.communicator.HaloUpdateRequest¶

Start an asynchronous halo update on a quantity.

Parameters:	quantity – the quantity to be updated n_points – how many halo points to update, starting from the interior
Returns:	an asynchronous request object with a .wait() method
Return type:	request

start_synchronize_vector_interfaces(x_quantity: fv3gfs.util.quantity.Quantity, y_quantity: fv3gfs.util.quantity.Quantity)¶

Synchronize shared points at the edges of a vector interface variable.

Sends the values on the south and west edges to overwrite the values on adjacent subtiles. Vector must be defined on the Arakawa C grid.

For interface variables, the edges of the tile are computed on both ranks bordering that edge. This routine copies values across those shared edges so that both ranks have the same value for that edge. It also handles any rotation of vector quantities needed to move data across the edge.

Parameters:	x_quantity – the x-component quantity to be synchronized y_quantity – the y-component quantity to be synchronized
Returns:	an asynchronous request object with a .wait() method
Return type:	request

start_vector_halo_update(x_quantity: fv3gfs.util.quantity.Quantity, y_quantity: fv3gfs.util.quantity.Quantity, n_points: int) → fv3gfs.util.communicator.HaloUpdateRequest¶

Start an asynchronous halo update of a horizontal vector quantity.

Assumes the x and y dimension indices are the same between the two quantities.

Parameters:	x_quantity – the x-component quantity to be halo updated y_quantity – the y-component quantity to be halo updated n_points – how many halo points to update, starting at the interior
Returns:	an asynchronous request object with a .wait() method
Return type:	request

synchronize_vector_interfaces(x_quantity: fv3gfs.util.quantity.Quantity, y_quantity: fv3gfs.util.quantity.Quantity)¶

Synchronize shared points at the edges of a vector interface variable.

Sends the values on the south and west edges to overwrite the values on adjacent subtiles. Vector must be defined on the Arakawa C grid.

For interface variables, the edges of the tile are computed on both ranks bordering that edge. This routine copies values across those shared edges so that both ranks have the same value for that edge. It also handles any rotation of vector quantities needed to move data across the edge.

Parameters:	x_quantity – the x-component quantity to be synchronized y_quantity – the y-component quantity to be synchronized

tile¶: communicator for within a tile

vector_halo_update(x_quantity: fv3gfs.util.quantity.Quantity, y_quantity: fv3gfs.util.quantity.Quantity, n_points: int)¶

Perform a halo update of a horizontal vector quantity.

Assumes the x and y dimension indices are the same between the two quantities.

Parameters:	x_quantity – the x-component quantity to be halo updated y_quantity – the y-component quantity to be halo updated n_points – how many halo points to update, starting at the interior

class fv3gfs.util.Communicator(comm, partitioner)¶

Bases: object

gather(send_quantity: fv3gfs.util.quantity.Quantity, recv_quantity: fv3gfs.util.quantity.Quantity = None) → Optional[fv3gfs.util.quantity.Quantity]¶

Transfer subtile regions of a full-tile quantity from each rank to the tile root rank.

Parameters:	send_quantity – quantity to send recv_quantity – if provided, assign received data into this Quantity (only used on the tile root rank)
Returns:	quantity if on root rank, otherwise None
Return type:	recv_quantity

gather_state(send_state=None, recv_state=None)¶

Transfer a state dictionary from subtile ranks to the tile root rank.

‘time’ is assumed to be the same on all ranks, and its value will be set to the value from the root rank.

Parameters:	send_state – the model state to be sent containing the subtile data recv_state – the pre-allocated state in which to recieve the full tile state. Only variables which are scattered will be written to.
Returns:	on the root rank, the state containing the entire tile
Return type:	recv_state

rank¶: rank of the current process within this communicator

scatter(send_quantity: Optional[fv3gfs.util.quantity.Quantity] = None, recv_quantity: Optional[fv3gfs.util.quantity.Quantity] = None) → fv3gfs.util.quantity.Quantity¶

Transfer subtile regions of a full-tile quantity from the tile root rank to all subtiles.

Parameters:	send_quantity – quantity to send, only required/used on the tile root rank recv_quantity – if provided, assign received data into this Quantity.
Returns:	recv_quantity

scatter_state(send_state=None, recv_state=None)¶

Transfer a state dictionary from the tile root rank to all subtiles.

Parameters:	send_state – the model state to be sent containing the entire tile, required only from the root rank recv_state – the pre-allocated state in which to recieve the scattered state. Only variables which are scattered will be written to.
Returns:	the state corresponding to this rank’s subdomain
Return type:	rank_state

class fv3gfs.util.HaloUpdateRequest(send_requests, recv_requests)¶

Bases: object

asynchronous request object for halo updates

wait()¶

fv3gfs.util.to_dataset(state)¶

fv3gfs.util.capture_stream(stream)¶

class fv3gfs.util.SubtileGridSizer(nx: int, ny: int, nz: int, n_halo: int, extra_dim_lengths: Dict[str, int])¶

Bases: fv3gfs.util.initialization.sizer.GridSizer

dim_extents¶

classmethod from_namelist(namelist: dict, tile_partitioner: fv3gfs.util.partitioner.TilePartitioner = None, tile_rank: int = 0)¶

Create a SubtileGridSizer from a Fortran namelist.

Parameters:	namelist – A namelist for the fv3gfs fortran model tile_partitioner (optional) – a partitioner to use for segmenting the tile. By default, a TilePartitioner is used. tile_rank (optional) – current rank on tile. Default is 0. Only matters if different ranks have different domain shapes. If tile_partitioner is a TilePartitioner, this argument does not matter.

classmethod from_tile_params(nx_tile: int, ny_tile: int, nz: int, n_halo: int, extra_dim_lengths: Dict[str, int], layout: Tuple[int, int], tile_partitioner: fv3gfs.util.partitioner.TilePartitioner = None, tile_rank: int = 0)¶

Create a SubtileGridSizer from parameters about the full tile.

Parameters:

nx_tile – number of x cell centers on the tile
ny_tile – number of y cell centers on the tile
nz – number of vertical levels
n_halo – number of halo points
extra_dim_lengths – lengths of any non-x/y/z dimensions, such as land or radiation dimensions
layout – (y, x) number of ranks along tile edges
tile_partitioner (optional) – partitioner object for the tile. By default, a TilePartitioner is created with the given layout
tile_rank (optional) – rank of this subtile.

get_extent(dims: Iterable[str]) → Tuple[int, ...]¶

get_origin(dims: Iterable[str]) → Tuple[int, ...]¶

get_shape(dims: Iterable[str]) → Tuple[int, ...]¶

class fv3gfs.util.GridSizer(nx: int, ny: int, nz: int, n_halo: int, extra_dim_lengths: Dict[str, int])¶

Bases: object

extra_dim_lengths = None¶: lengths of any non-x/y/z dimensions, such as land or radiation dimensions

get_extent(dims: Sequence[str]) → Tuple[int, ...]¶

get_origin(dims: Sequence[str]) → Tuple[int, ...]¶

get_shape(dims: Sequence[str]) → Tuple[int, ...]¶

n_halo = None¶: number of horizontal halo points for produced arrays

nx = None¶: length of the x compute dimension for produced arrays

ny = None¶: length of the y compute dimension for produced arrays

nz = None¶: length of the z compute dimension for produced arrays

class fv3gfs.util.QuantityFactory(sizer: fv3gfs.util.initialization.sizer.SubtileGridSizer, numpy)¶

Bases: object

empty(dims: Sequence[str], units: str, dtype: type = <class 'float'>)¶

classmethod from_backend(sizer: fv3gfs.util.initialization.sizer.SubtileGridSizer, backend: str)¶

Initialize a QuantityFactory to use a specific gt4py backend.

Parameters:	sizer – object which determines array sizes backend – gt4py backend

ones(dims: Sequence[str], units: str, dtype: type = <class 'float'>)¶

zeros(dims: Sequence[str], units: str, dtype: type = <class 'float'>)¶

fv3gfs.util.array_buffer(allocator: fv3gfs.util.types.Allocator, shape: Iterable[int], dtype: type)¶

A context manager providing a contiguous array, which may be re-used between calls.

Parameters:

allocator – a function with the same signature as numpy.zeros which returns an ndarray
shape – the shape of the desired array
dtype – the dtype of the desired array

Yields:

buffer_array –

an ndarray created according to the specification in the args.: May be retained and re-used in subsequent calls.

fv3gfs.util.send_buffer(allocator: Callable, array: numpy.ndarray, timer: Optional[fv3gfs.util._timing.Timer] = None)¶

A context manager ensuring that array is contiguous in a context where it is being sent as data, copying into a recycled buffer array if necessary.

Parameters:

allocator – a function behaving like numpy.empty
array – a possibly non-contiguous array for which to provide a buffer
timer – object to accumulate timings for “pack”

Yields:

buffer_array –

if array is non-contiguous, a contiguous buffer array containing: the data from array. Otherwise, yields array.

fv3gfs.util.recv_buffer(allocator: Callable, array: numpy.ndarray, timer: Optional[fv3gfs.util._timing.Timer] = None)¶

A context manager ensuring that array is contiguous in a context where it is being used to receive data, using a recycled buffer array and then copying the result into array if necessary.

Parameters:

allocator – a function behaving like numpy.empty
array – a possibly non-contiguous array for which to provide a buffer
timer – object to accumulate timings for “unpack”

Yields:

buffer_array –

if array is non-contiguous, a contiguous buffer array which is: copied into array when the context is exited. Otherwise, yields array.

fv3gfs.util.fill_scalar_corners(quantity: fv3gfs.util.quantity.Quantity, direction: typing_extensions.Literal['x', 'y'][x, y], tile_partitioner: fv3gfs.util.partitioner.TilePartitioner, rank: int, n_halo: int)¶

At the corners of tile faces, copy data from halo edges into halo corners to allow stencils to be translated along those edges in a computationally-relevant way.

The quantity is modified in-place.

Parameters:

quantity – the quantity to modify, whose first two dimensions must be along the x and y directions, respectively
direction – the direction along which we want to enable stencils to compute. For example, calling with “x” would allow a stencil with length > 1 along the x-direction to be convolved with Quantity. Note it is not possible to use corner filling to convolve with stencils having length > 1 along both x and y dimensions.
tile_partitioner – object to determine tile positions of ranks
rank – rank on which the quantity exists
n_halo – number of halo points to fill