credit.postblock.geopotential

Classes

GeopotentialDiagnostic

GeopotentialDiagnostic is a neural network module used for computing geopotential

Functions

pressure_on_interfaces(surface_pressure, model_a_half, ...)	Calculate pressure on the interfaces of atmospheric hybrid sigma-pressure model levels.
geopotential(surface_geopotential, surface_pressure, ...)	Calculate geopotential (m**2 s**-2) from hybrid sigma-pressure atmospheric level data.

Module Contents

credit.postblock.geopotential.pressure_on_interfaces(surface_pressure: torch.Tensor, model_a_half: torch.Tensor, model_b_half: torch.Tensor, model_top_pressure: float = 0.57)

Calculate pressure on the interfaces of atmospheric hybrid sigma-pressure model levels. Conversion is pressure_3d = a + b * SP. The a and b coefficients are defined at the interfaces of the levels rather than at the level centers where the mass variables are defined.

Parameters:

• surface_pressure (torch.Tensor) – (time, latitude, longitude) or (latitude, longitude) grid in units of Pa.

• model_a_half (torch.Tensor) – coefficient a at each model level interface in units of Pa.

• model_b_half (torch.Tensor) – coefficient b at each model level interface, which is unitless.

• model_top_pressure (float) – pressure at model top (default 0.57 Pa).

Returns:

Pressure on each model level interface.

credit.postblock.geopotential.geopotential(surface_geopotential: torch.Tensor, surface_pressure: torch.Tensor, temperature: torch.Tensor, specific_humidity: torch.Tensor, model_a_half: torch.Tensor, model_b_half: torch.Tensor, flip_vertical: bool = True)

Calculate geopotential (m**2 s**-2) from hybrid sigma-pressure atmospheric level data.

This calculation assumes that the 3D fields are oriented in the vertical/levels dimension such that pressure increases with increasing index value (top of atmosphere to surface).

Parameters:

• surface_geopotential (torch.Tensor) – surface geopotential (m**2 s**-2)

• surface_pressure (torch.Tensor) – surface pressure (Pa)

• temperature (torch.Tensor) – temperature (K)

• specific_humidity (torch.Tensor) – specific humidity (kg kg**-1)

• model_a_half (torch.Tensor) – pressure component a at each model level interface (Pa)

• model_b_half (torch.Tensor) – sigma component b at each model level interface (unitless)

• flip_vertical (bool) – whether to flip the vertical dimension of the input tensors. Default True.

Returns:

geopotential on each model level (m**2 s**-2)

Return type:

geopotential (torch.Tensor)

class credit.postblock.geopotential.GeopotentialDiagnostic(output_name: str = 'ARCO_ERA5/derived_diagnostic/3d/geopotential', dataset_name: str = 'ARCO_ERA5', chunk_size: int = 1000, data_keys: Iterable[str] = ('prediction', 'target'), surface_geopotential_var: str = 'ARCO_ERA5/static/2d/geopotential_at_surface', surface_pressure_var: str = 'ARCO_ERA5/prognostic/2d/surface_pressure', temperature_var: str = 'ARCO_ERA5/prognostic/3d/temperature', specific_humidity_var: str = 'ARCO_ERA5/prognostic/3d/specific_humidity', flip_vertical: bool = True, level_info_file: str = 'ERA5_Lev_Info.nc', model_a_half_var: str = 'a_half', model_b_half_var: str = 'b_half', static_source_key: str = 'ic_raw', levels: list[int] | None = None)

Bases: torch.nn.Module

GeopotentialDiagnostic is a neural network module used for computing geopotential diagnostics using multi-dimensional input data.

This class processes geophysical variables such as surface geopotential, surface pressure, temperature, and specific humidity to calculate geopotential fields. The input data is expected to conform to a specific format, and the class makes use of auxiliary metadata files that describe model-specific level information.

output_name

The key used in the dataset to store the computed geopotential diagnostic output.

Type:

str

dataset_name

The name of the dataset from which input variables will be retrieved.

Type:

str

chunk_size

The chunk size used for vectorized computations to optimize memory usage during processing.

Type:

int

data_keys

The keys in the input data dictionary that will be processed (e.g., “prediction”, “target”).

Type:

Iterable[str]

surface_geopotential_var

The key for the surface geopotential variable in the dataset.

Type:

str

surface_pressure_var

The key for the surface pressure variable in the dataset.

Type:

str

temperature_var

The key for the temperature variable in the dataset.

Type:

str

specific_humidity_var

The key for the specific humidity variable in the dataset.

Type:

str

flip_vertical

Whether to flip the vertical dimension of the input tensors. Default True

Type:

bool

level_info_file

The filename of the auxiliary metadata file that stores information about model levels.

Type:

str

model_a_half_var

The variable name for the a (pressure) hybrid sigma-pressure coefficient in the level information file.

Type:

str

model_b_half_var

The variable name for the b (sigma) hybrid sigma-pressure coefficient parameter in the level information file.

Type:

str

output_name = 'ARCO_ERA5/derived_diagnostic/3d/geopotential'

dataset_name = 'ARCO_ERA5'

chunk_size = 1000

data_keys = ('prediction', 'target')

surface_geopotential_var = 'ARCO_ERA5/static/2d/geopotential_at_surface'

surface_pressure_var = 'ARCO_ERA5/prognostic/2d/surface_pressure'

temperature_var = 'ARCO_ERA5/prognostic/3d/temperature'

specific_humidity_var = 'ARCO_ERA5/prognostic/3d/specific_humidity'

flip_vertical = True

level_info_file

model_a_half_var = 'a_half'

model_b_half_var = 'b_half'

static_source_key = 'ic_raw'

levels = None

forward(data_dict: dict)

Processes a dictionary of input data, rearranges dimensions, computes derived quantities using a custom function geopotential, and updates the data dictionary with the results.

Parameters:

data_dict (dict) – Input dictionary containing data corresponding to various data types. The data for each type is expected to be organized into specified attributes (e.g., temperature, specific humidity).

Returns:

Updated data dictionary, where new computed fields are added to the relevant dataset, preserving the original structure.

Return type:

dict

Raises:

ValueError – If any required data type is not found in the input data_dict.