ObsGridMapping#

class earth2studio.utils.obs.ObsGridMapping(grid_variables, grid_lat, grid_lon, device=None)[source]#

Maps point observations to and from a 2-D model grid.

Supports regular (1-D lat/lon axes) and irregular (2-D lat/lon) grids. CuPy is used automatically when available and the target device is a GPU.

Parameters:

  • grid_variables (np.ndarray) – Ordered array of variable names that define the channel dimension.

  • grid_lat (np.ndarray or torch.Tensor) – Latitude coordinates. Shape (H,) for a regular grid or (H, W) for an irregular grid.

  • grid_lon (np.ndarray or torch.Tensor) – Longitude coordinates. Shape (W,) for a regular grid or (H, W) for an irregular grid. Values are expected in [0, 360).

  • device (torch.device, str, or None, optional) – Target device. Inferred from grid_lat when it is a torch.Tensor; defaults to CPU otherwise.

Raises:

ValueError – If grid_lat is neither 1-D nor 2-D.

grid_to_obs(x, obs_c, obs_i, obs_j, method='linear')[source]#

Sample grid values at observation locations (differentiable).

Parameters:

  • x (torch.Tensor) – Grid tensor of shape (C, H, W).

  • obs_c (torch.Tensor) – Channel index for each observation.

  • obs_i (torch.Tensor) – Fractional row index for each observation.

  • obs_j (torch.Tensor) – Fractional column index for each observation.

  • method ({"nearest", "linear"}, optional) – Interpolation method. "linear" (default) performs bilinear interpolation and supports gradient backpropagation.

Returns:

Sampled values, shape (N,).

Return type:

torch.Tensor

Raises:

ValueError – If method is not "nearest" or "linear".

obs_coords(
obs_var,
obs_lat,
obs_lon,
remove_out_of_bounds=True,
)[source]#

Convert observation locations to fractional grid indices.

Parameters:

  • obs_var (np.ndarray) – Variable name for each observation.

  • obs_lat (np.ndarray) – Latitude of each observation in degrees.

  • obs_lon (np.ndarray) – Longitude of each observation in degrees, in [0, 360).

  • remove_out_of_bounds (bool, optional) – When True (default), observations outside the grid or with an unknown variable are removed from the returned arrays. The full boolean mask is always returned as the fourth element.

Returns:

  • obs_c (torch.Tensor) – Channel index for each (retained) observation.

  • obs_i (torch.Tensor) – Fractional row index in the grid for each observation.

  • obs_j (torch.Tensor) – Fractional column index in the grid for each observation.

  • obs_in_bounds (torch.Tensor) – Boolean mask of length N indicating which of the original observations are in-bounds and have a known variable.

Return type:

tuple[Tensor, Tensor, Tensor, Tensor]

obs_to_grid(
obs,
variables=None,
request_time=None,
time_tolerance=None,
return_empty_grid=False,
)[source]#

Bin point observations onto the model grid.

Observations are filtered to variables and to those within time_tolerance of request_time. Multiple observations that fall in the same grid cell are averaged.

Parameters:

  • obs (pd.DataFrame or None) – Observation table with columns variable, lat, lon, observation. If request_time is used, obs must also have a time column.

  • variables (list[str], optional) – Variable names to retain; others are discarded. None (default) uses all variables.

  • request_time (np.datetime64, optional) – Target valid time for the observations. None (default) skips time filtering.

  • time_tolerance (np.timedelta64, tuple[np.timedelta64, np.timedelta64], optional) – Time window used to filter observations around request_time. A single np.timedelta64 creates a symmetric window [-time_tolerance, +time_tolerance]. A 2-tuple (lower, upper) is passed directly to filter_time_range as an asymmetric window. None (default) skips time filtering.

  • return_empty_grid (bool, optional) – When True, return zero-filled tensors instead of (None, None) when no observations are gridded. Defaults to False.

Returns:

  • y_obs (torch.Tensor or None) – Gridded observations, shape (1, C, H, W), or None when no observations remain and return_empty_grid is False.

  • mask (torch.Tensor or None) – Observation count per cell (float), same shape as y_obs, or None under the same condition as above.

Return type:

tuple[Tensor, Tensor] | tuple[None, None]