credit.physics_core =================== .. py:module:: credit.physics_core .. autoapi-nested-parse:: Tools for physics-based constraints and derivations for CREDIT models -------------------------------------------------------------------------- Content: - physics_pressure_level - physics_hybrid_sigma_level Reference: - https://journals.ametsoc.org/view/journals/clim/34/10/JCLI-D-20-0676.1.xml - https://doi.org/10.1175/JCLI-D-13-00018.1 - https://github.com/ai2cm/ace/tree/main/fme/fme/core Classes ------- .. autoapisummary:: credit.physics_core.ModelLevelPressures credit.physics_core.physics_pressure_level credit.physics_core.physics_hybrid_sigma_level Functions --------- .. autoapisummary:: credit.physics_core.compute_density credit.physics_core.compute_virtual_temperature Module Contents --------------- .. py:function:: compute_density(pressure, temperature, specific_humidity) compute density given pressure (Pa), temperature (K), and specific humidity (kg/kg) .. py:function:: compute_virtual_temperature(temperature, specific_humidity) ref: metpy .. py:class:: ModelLevelPressures(a_vals, b_vals, plev_dim=1) Bases: :py:obj:`torch.nn.Module` compute pressure levels given SP with (only compatible with torch) SP, a_vals with same units. a_vals, b_vals, with size levels at dimension plev_dim and all other dims with size 1, e.g. plev_dim = 1; a_vals.shape = (1, levels, 1, 1, 1) matching sp with size 1 at dimension plev_dim e.g. (b, 1, t, lat, lon) .. py:attribute:: plev_dim :value: 1 .. py:attribute:: is_fully_initialized :value: False .. py:method:: compute_p(sp) .. py:method:: compute_hlevs(plevs) .. py:method:: compute_mlev_thickness(sp) .. py:class:: physics_pressure_level(lon: torch.Tensor, lat: torch.Tensor, upper_air_pressure: torch.Tensor, midpoint: bool = False) Pressure level physics All inputs must be in the same torch device. Full order of dimensions: (batch, time, level, latitude, longitude) .. py:attribute:: lon .. py:attribute:: lat .. py:attribute:: upper_air_pressure .. py:attribute:: pressure_thickness .. py:attribute:: area .. py:method:: pressure_integral_midpoint(q_mid: torch.Tensor) -> torch.Tensor Compute the pressure level integral of a given quantity; assuming its mid point values are pre-computed :param q_mid: the quantity with dims of (batch_size, time, level-1, latitude, longitude) :returns: Pressure level integrals of q .. py:method:: pressure_integral_midpoint_sliced(q_mid: torch.Tensor, ind_start: int, ind_end: int) -> torch.Tensor As in `pressure_integral_midpoint`, but supports pressure level indexing, so it can calculate integrals of a subset of levels .. py:method:: pressure_integral_trapz(q: torch.Tensor) -> torch.Tensor Compute the pressure level integral of a given quantity using the trapezoidal rule. :param q: the quantity with dims of (batch_size, time, level, latitude, longitude) :returns: Pressure level integrals of q .. py:method:: pressure_integral_trapz_sliced(q: torch.Tensor, ind_start: int, ind_end: int) -> torch.Tensor As in `pressure_integral_trapz`, but supports pressure level indexing, so it can calculate integrals of a subset of levels .. py:method:: weighted_sum(q: torch.Tensor, axis: Dict[tuple, None] = None, keepdims: bool = False) -> torch.Tensor Compute the weighted sum of a given quantity for PyTorch tensors. :param q: the quantity to be summed (PyTorch tensor) :param axis: dims to compute the sum (can be int or tuple of ints) :param keepdims: whether to keep the reduced dimensions or not :returns: Weighted sum (PyTorch tensor) .. py:method:: total_dry_air_mass(q: torch.Tensor) -> torch.Tensor Compute the total mass of dry air over the entire globe [kg] .. py:method:: total_column_water(q: torch.Tensor) -> torch.Tensor Compute total column water (TCW) per air column [kg/m2] .. py:class:: physics_hybrid_sigma_level(lon: torch.Tensor, lat: torch.Tensor, coef_a: torch.Tensor, coef_b: torch.Tensor, midpoint: bool = False) Hybrid sigma-pressure level physics .. attribute:: lon Longitude in degrees. :type: torch.Tensor .. attribute:: lat Latitude in degrees. :type: torch.Tensor .. attribute:: surface_pressure Surface pressure in Pa. :type: torch.Tensor .. attribute:: coef_a Hybrid sigma-pressure coefficient 'a' [Pa]. :type: torch.Tensor .. attribute:: coef_b Hybrid sigma-pressure coefficient 'b' [unitless]. :type: torch.Tensor .. attribute:: area Area of grid cells [m^2]. :type: torch.Tensor .. attribute:: integral Vertical integration method (midpoint or trapezoidal). :type: function .. py:attribute:: lon .. py:attribute:: lat .. py:attribute:: coef_a .. py:attribute:: coef_b .. py:attribute:: area .. py:method:: pressure_integral_midpoint(q_mid: torch.Tensor, surface_pressure: torch.Tensor) -> torch.Tensor Compute the pressure level integral of a given quantity; assuming its mid-point values are pre-computed. :param q_mid: The quantity with dims of (batch, level-1, time, latitude, longitude) :param surface_pressure: Surface pressure in Pa (batch, time, latitude, longitude). :returns: Pressure level integrals of q .. py:method:: pressure_integral_midpoint_sliced(q_mid: torch.Tensor, surface_pressure: torch.Tensor, ind_start: int, ind_end: int) -> torch.Tensor As in `pressure_integral_midpoint`, but supports pressure level indexing, so it can calculate integrals of a subset of levels. .. py:method:: pressure_integral_trapz(q: torch.Tensor, surface_pressure: torch.Tensor) -> torch.Tensor Compute the pressure level integral of a given quantity using the trapezoidal rule. :param q: The quantity with dims of (batch, level, time, latitude, longitude) :returns: Pressure level integrals of q .. py:method:: pressure_integral_trapz_sliced(q: torch.Tensor, surface_pressure: torch.Tensor, ind_start: int, ind_end: int) -> torch.Tensor As in `pressure_integral_trapz`, but supports pressure level indexing, so it can calculate integrals of a subset of levels. .. py:method:: weighted_sum(q: torch.Tensor, axis: Dict[tuple, None] = None, keepdims: bool = False) -> torch.Tensor Compute the weighted sum of a given quantity for PyTorch tensors. :param data: the quantity to be summed (PyTorch tensor) :param axis: dims to compute the sum (can be int or tuple of ints) :param keepdims: whether to keep the reduced dimensions or not :returns: Weighted sum (PyTorch tensor) .. py:method:: total_dry_air_mass(q: torch.Tensor, surface_pressure: torch.Tensor) -> torch.Tensor Compute the total mass of dry air over the entire globe [kg] .. py:method:: total_column_water(q: torch.Tensor, surface_pressure: torch.Tensor) -> torch.Tensor Compute total column water (TCW) per air column [kg/m2]