Skip to content

DLESyMv0_ISCCP_ERA5Precip

GlobalCM202440 GBNVIDIAPyTorch

Import path: earth2studio.models.dx.DLESyMv0_ISCCP_ERA5Precip

View source on GitHub View install commands

Documentation

Bases: Module, AutoModelMixin

Precipitation diagnostic for the DLESyMv0_ISCCP_ERA5 climate model.

A HEALPixUNet diagnostic from the AtmosSci-DLESM/DLESyM repository that predicts 6-hourly accumulated precipitation (tp06) from the full coupled atmosphere/ocean state. It is designed to be chained off [DLESyMv0_ISCCP_ERA5][earth2studio.models.px.DLESyMv0_ISCCP_ERA5] (or its lat/lon variant).

The model takes 2 consecutive history timesteps of 10 variables (the 9 atmospheric variables plus sst) on a HEALPix nside=64 grid and predicts tp06 at the last input timestep (i.e., precipitation accumulated over the 6 h prior to t=0).

When use_ttr=True (default), the wrapper accepts ERA5 ttr in place of rlut and applies the same per-doy moment-matching TTR -> OLR transform as [DLESyMv0_ISCCP_ERA5][earth2studio.models.px.dlesym_v0_isccp_era5.DLESyMv0_ISCCP_ERA5] before the forward pass. This allows the diagnostic to run standalone from an ERA5 initial condition. When use_ttr=False, supply pre-transformed rlut directly -- use this when chaining off [DLESyMv0_ISCCP_ERA5][earth2studio.models.px.DLESyMv0_ISCCP_ERA5] output, which is already in OLR space.

Parameters:

  • core_model (Module) –

    The wrapped physicsnemo.models.dlwp_healpix.HEALPixUNet.

  • hpx_lat (ndarray) –

    HEALPix latitude grid, shape (12, nside, nside).

  • hpx_lon (ndarray) –

    HEALPix longitude grid, shape (12, nside, nside).

  • nside (int) –

    HEALPix nside.

  • center (ndarray) –

    Per-variable input means, shape (1, 1, 1, V, 1, 1, 1) ordered by variables.

  • scale (ndarray) –

    Per-variable input stds, same shape as center.

  • constants (ndarray) –

    Constant fields (e.g. land-sea mask, topography), shape (12, n_constants, nside, nside).

  • input_times (ndarray) –

    Input lead times, shape (input_time_dim,). e.g. np.array([-6, 0], dtype='timedelta64[h]').

  • variables (list[str]) –

    Input variable names in channel order. Must contain rlut (the model variable-space name); input_coords swaps it to ttr when use_ttr=True.

  • output_variable (str, default: 'tp06' ) –

    Output variable name, typically "tp06".

  • log_epsilon (float | None, default: 1e-08 ) –

    If non-None, denormalize as exp(out + log(eps)) - eps to invert the upstream log-transform applied to precipitation during training. Set to None to disable.

  • use_ttr (bool, default: True ) –

    Accept ERA5 ttr instead of rlut and run the TTR -> OLR transform internally. Defaults to True.

  • ttr_clim_mean (ndarray | None, default: None ) –

    Per-doy climatology arrays of shape (D, F, H, W). Required when use_ttr=True.

  • ttr_clim_std (ndarray | None, default: None ) –

    Per-doy climatology arrays of shape (D, F, H, W). Required when use_ttr=True.

  • olr_clim_mean (ndarray | None, default: None ) –

    Per-doy climatology arrays of shape (D, F, H, W). Required when use_ttr=True.

  • olr_clim_std (ndarray | None, default: None ) –

    Per-doy climatology arrays of shape (D, F, H, W). Required when use_ttr=True.

  • olr_floor (float, default: 0.0 ) –

    Lower bound applied to the transformed OLR. Defaults to 0.0.

Note

For more information see:

__call__ 

__call__(
    x: Tensor, coords: CoordSystem
) -> tuple[Tensor, CoordSystem]

Run the precip diagnostic forward.

Parameters:

  • x (Tensor) –

    Input of shape (B, T, LT, V, F, H, W) with LT = input_time_dim history timesteps and V = len(variables) channels on HEALPix.

  • coords (CoordSystem) –

    Input coordinates.

Returns:

  • tuple[Tensor, CoordSystem] –

    Output tensor of shape (B, T, 1, 1, F, H, W) and the corresponding output coords with variable = [tp06].

load_default_package classmethod 

load_default_package() -> Package

Default DLESyMv0_ISCCP_ERA5 precip package on HuggingFace.

load_model classmethod 

load_model(
    package: Package, use_ttr: bool = True
) -> DiagnosticModel

Load the DLESyMv0_ISCCP_ERA5 precip diagnostic from a package.

Parameters:

  • package (Package) –

    Model package containing config.yaml, the precip .mdlus checkpoint, the HEALPix lat/lon, the constant fields, and (when use_ttr=True) the TTR/OLR climatology netCDFs (era5_ttr_doy_stats_hpx64.nc and isccp_olr_doy_stats_hpx64.nc).

  • use_ttr (bool, default: True ) –

    See DLESyMv0_ISCCP_ERA5Precip. Defaults to True.

Returns:

  • DiagnosticModel –

    Loaded DLESyMv0_ISCCP_ERA5Precip instance.

Examples using earth2studio.models.dx.DLESyMv0_ISCCP_ERA5Precip

Preview of Running the DLESyM ISCCP-ERA5 ModelRunning the DLESyM ISCCP-ERA5 ModelCoupled inference with the upstream DLESyM checkpoints, including precipitation.