credit.postblock.geopotential ============================= .. py:module:: credit.postblock.geopotential Classes ------- .. autoapisummary:: credit.postblock.geopotential.GeopotentialDiagnostic Functions --------- .. autoapisummary:: credit.postblock.geopotential.pressure_on_interfaces credit.postblock.geopotential.geopotential Module Contents --------------- .. py:function:: 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. :param surface_pressure: (time, latitude, longitude) or (latitude, longitude) grid in units of Pa. :type surface_pressure: torch.Tensor :param model_a_half: coefficient a at each model level interface in units of Pa. :type model_a_half: torch.Tensor :param model_b_half: coefficient b at each model level interface, which is unitless. :type model_b_half: torch.Tensor :param model_top_pressure: pressure at model top (default 0.57 Pa). :type model_top_pressure: float :returns: Pressure on each model level interface. .. py:function:: 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). :param surface_geopotential: surface geopotential (m**2 s**-2) :type surface_geopotential: torch.Tensor :param surface_pressure: surface pressure (Pa) :type surface_pressure: torch.Tensor :param temperature: temperature (K) :type temperature: torch.Tensor :param specific_humidity: specific humidity (kg kg**-1) :type specific_humidity: torch.Tensor :param model_a_half: pressure component a at each model level interface (Pa) :type model_a_half: torch.Tensor :param model_b_half: sigma component b at each model level interface (unitless) :type model_b_half: torch.Tensor :param flip_vertical: whether to flip the vertical dimension of the input tensors. Default True. :type flip_vertical: bool :returns: geopotential on each model level (m**2 s**-2) :rtype: geopotential (torch.Tensor) .. py:class:: 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: :py:obj:`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. .. attribute:: output_name The key used in the dataset to store the computed geopotential diagnostic output. :type: str .. attribute:: dataset_name The name of the dataset from which input variables will be retrieved. :type: str .. attribute:: chunk_size The chunk size used for vectorized computations to optimize memory usage during processing. :type: int .. attribute:: data_keys The keys in the input data dictionary that will be processed (e.g., "prediction", "target"). :type: Iterable[str] .. attribute:: surface_geopotential_var The key for the surface geopotential variable in the dataset. :type: str .. attribute:: surface_pressure_var The key for the surface pressure variable in the dataset. :type: str .. attribute:: temperature_var The key for the temperature variable in the dataset. :type: str .. attribute:: specific_humidity_var The key for the specific humidity variable in the dataset. :type: str .. attribute:: flip_vertical Whether to flip the vertical dimension of the input tensors. Default True :type: bool .. attribute:: level_info_file The filename of the auxiliary metadata file that stores information about model levels. :type: str .. attribute:: model_a_half_var The variable name for the `a` (pressure) hybrid sigma-pressure coefficient in the level information file. :type: str .. attribute:: model_b_half_var The variable name for the `b` (sigma) hybrid sigma-pressure coefficient parameter in the level information file. :type: str .. py:attribute:: output_name :value: 'ARCO_ERA5/derived_diagnostic/3d/geopotential' .. py:attribute:: dataset_name :value: 'ARCO_ERA5' .. py:attribute:: chunk_size :value: 1000 .. py:attribute:: data_keys :value: ('prediction', 'target') .. py:attribute:: surface_geopotential_var :value: 'ARCO_ERA5/static/2d/geopotential_at_surface' .. py:attribute:: surface_pressure_var :value: 'ARCO_ERA5/prognostic/2d/surface_pressure' .. py:attribute:: temperature_var :value: 'ARCO_ERA5/prognostic/3d/temperature' .. py:attribute:: specific_humidity_var :value: 'ARCO_ERA5/prognostic/3d/specific_humidity' .. py:attribute:: flip_vertical :value: True .. py:attribute:: level_info_file .. py:attribute:: model_a_half_var :value: 'a_half' .. py:attribute:: model_b_half_var :value: 'b_half' .. py:attribute:: static_source_key :value: 'ic_raw' .. py:attribute:: levels :value: None .. py:method:: 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. :param data_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). :type data_dict: dict :returns: Updated data dictionary, where new computed fields are added to the relevant dataset, preserving the original structure. :rtype: dict :raises ValueError: If any required data type is not found in the input `data_dict`.