credit.postblock.hybrid_interp

credit.postblock.hybrid_interp#

Interpolates 3D variables from one set of hybrid sigma-pressure levels to another, following credit.interp.interp_hybrid_to_hybrid_levels / credit.gefs.interpolate_vertical_levels: source and destination pressure columns are both built from the same surface pressure field (p = a + b * SP), and each variable is interpolated linearly in log(pressure), column by column, parallelized with torch.vmap. Values outside the source pressure range use constant extrapolation (np.interp behavior).

The primary use case is inference with a model trained on one vertical grid (e.g. ERA5 hybrid levels) driven by initial conditions from another (e.g. GFS): run this as a preblock (credit.preblock.hybrid_interp.HybridLevelInterp) to move the IC onto the model’s levels, or as a postblock to move model output onto another grid. This module holds the shared engine and the postblock; the preblock is a thin wrapper around the same engine.

Classes#

HybridLevelInterpPost

Postblock that interpolates 3D variables between hybrid sigma-pressure level sets.

Functions#

interp_column_hybrid_to_hybrid(→ torch.Tensor)

Interpolate one column between two hybrid sigma-pressure level sets.

Module Contents#

credit.postblock.hybrid_interp.interp_column_hybrid_to_hybrid(fields: torch.Tensor, surface_pressure: torch.Tensor, source_a: torch.Tensor, source_b: torch.Tensor, dest_a: torch.Tensor, dest_b: torch.Tensor) → torch.Tensor#

Interpolate one column between two hybrid sigma-pressure level sets.

Both pressure columns are built from the same surface pressure (p = a + b * SP); interpolation is linear in log(pressure) with constant extrapolation outside the source range. Source levels must be ordered top-of-atmosphere to surface (pressure increasing with index); destination levels may be in any order. Designed to be wrapped in torch.vmap over columns: fields and surface_pressure carry a batch dimension in the vmapped call, the coefficients do not.

Parameters:

  • fields – variables to interpolate, shape (n_fields, n_source_levels).

  • surface_pressure – surface pressure in Pa, scalar tensor.

  • source_a – source a coefficient at level midpoints in Pa, shape (n_source_levels,).

  • source_b – source b coefficient at level midpoints (unitless), shape (n_source_levels,).

  • dest_a – destination a coefficient at level midpoints in Pa, shape (n_dest_levels,).

  • dest_b – destination b coefficient at level midpoints (unitless), shape (n_dest_levels,).

Returns:

Interpolated fields, shape (n_fields, n_dest_levels).

class credit.postblock.hybrid_interp.HybridLevelInterpPost(key: str = 'y_processed', **engine_kwargs)#

Bases: credit.postblock.base.BasePostblock

Postblock that interpolates 3D variables between hybrid sigma-pressure level sets.

Follows the same batch-dict protocol as Reconstruct and BridgeScalerTransform: it operates on the nested output dict at batch_dict[key] (default "y_processed"), which has the form {source: {var_key: tensor}}, replacing each variable in variables with its interpolated (B, n_dest_levels, n_time, H, W) counterpart in place. A preblock with the same engine and arguments is available as credit.preblock.hybrid_interp.HybridLevelInterp for regridding initial conditions (e.g. GFS → ERA5 levels) before the model runs.

Coefficient files: bare filenames resolve to credit.metadata; paths with directories (env vars allowed) are used as-is. Coefficients on level interfaces (*_on_interfaces: true, e.g. ERA5 a_half/b_half or CAM hyai/hybi) are averaged to midpoints as in credit.interp.create_pressure_grid; set false for midpoint coefficients (e.g. a_model/b_model). If the a and b variable names are equal and the variable is 2D (the GFS vcoord convention), row 0 is a and row 1 is b.

Config example (GFS-level output → ERA5 levels):

type: "hybrid_level_interp"
args:
    variables:
        - "GFS/prognostic/3d/temperature"
        - "GFS/prognostic/3d/specific_humidity"
    surface_pressure_var: "GFS/prognostic/2d/surface_pressure"
    source_level_info_file: "/path/to/gfs_ctrl.nc"
    source_a_var: "vcoord"
    source_b_var: "vcoord"
    dest_level_info_file: "ERA5_Lev_Info.nc"
Parameters:

  • variables – 3D var_keys to interpolate; variables absent from the batch are skipped silently.

  • surface_pressure_var – variable name for surface pressure (Pa).

  • source_level_info_file – file with the levels the data is currently on.

  • dest_level_info_file – file with the levels to interpolate to.

  • source_a_var – name of the source a coefficient variable (Pa).

  • source_b_var – name of the source b coefficient variable (unitless).

  • source_on_interfaces – whether source coefficients are on level interfaces.

  • source_levels – optional subset of 1-based source midpoint level numbers.

  • dest_a_var – name of the destination a coefficient variable (Pa).

  • dest_b_var – name of the destination b coefficient variable (unitless).

  • dest_on_interfaces – whether destination coefficients are on level interfaces.

  • dest_levels – optional subset of 1-based destination midpoint level numbers.

  • chunk_size – vmap chunk size to bound memory usage.

  • key – entry in batch_dict holding the nested output dict written by Reconstruct (default: "y_processed").

key = 'y_processed'#
engine#
forward(batch_dict: dict) → dict#