Data Utilities¶

def extract_wet_mask(
    data: xr.Dataset,
    prognostic_var_names: PrognosticVarNames,
    boundary_var_names: BoundaryVarNames,
) -> Masks:
    """A mask for where the oceans are. Water is wet."""
    data_ = flatten_masks(data)
    wet_inp_np = np.stack(
        [_mask_array_for_data_var(data_, var_name) for var_name in prognostic_var_names]
    )
    boundary_mask_vars = [
        _mask_var_for_data_var(data_, var_name) for var_name in boundary_var_names
    ]
    boundary_masks = [
        _mask_array_for_data_var(data_, var_name) for var_name in boundary_var_names
    ]
    wet_surface_mask_np = (
        np.stack(boundary_masks)
        if len(set(boundary_mask_vars)) > 1
        else boundary_masks[0]
    )

    wet_inp = torch.from_numpy(wet_inp_np)
    wet_surface = torch.from_numpy(wet_surface_mask_np)
    return Masks(wet_inp.bool(), wet_surface.bool())

def flatten_masks(
    data: xr.Dataset,
) -> xr.Dataset:
    """Adds level-wise mask variables from the stacked wet mask."""
    data_ = data.copy()
    if f"{_CANONICAL_MASK_PREFIX}0" not in data_.variables:
        assert _STACKED_WET_MASK_VAR in data_.variables, (
            "Wet mask cannot be constructed without "
            "either the wetmask variable or the level-wise masks"
        )

        wet_mask = data_[_STACKED_WET_MASK_VAR]
        for i in range(wet_mask.sizes["lev"]):
            data_[f"{_CANONICAL_MASK_PREFIX}{i}"] = wet_mask.isel(lev=i)

        data_ = data_.drop_vars(_STACKED_WET_MASK_VAR)

    return data_

def unflatten_masks(
    data: xr.Dataset,
    num_levels: int,
) -> xr.Dataset:
    """Adds a stacked wet mask `xarray.DataArray` from level-wise mask variables."""
    data_ = data.copy()
    mask_vars = [f"{_CANONICAL_MASK_PREFIX}{i}" for i in range(num_levels)]
    if _STACKED_WET_MASK_VAR not in data_.variables:
        assert mask_vars[0] in data_.variables, "Wet mask must have masks as data vars!"

        wetmask = data_[mask_vars].to_array(dim="lev", name=_STACKED_WET_MASK_VAR)

        lev = data_.coords.get("lev", np.arange(len(mask_vars)))
        data_[_STACKED_WET_MASK_VAR] = wetmask.assign_coords(lev=lev)
        data_ = data_.drop_vars(mask_vars)

    return data_

def spherical_area(data: xr.Dataset) -> Grid:
    """
    Compute real grid cell areas on a spherical Earth.

    Uses the spherical geometry formula:
    A = R² × Δλ × (sin(φ₂) - sin(φ₁))

    where:
    - R is Earth's radius (6371 km)
    - Δλ is the longitude spacing in radians
    - φ₁, φ₂ are the latitude bounds of the cell in radians

    Args:
        data: Dataset containing lat/lon coordinates

    Returns:
        Grid cell areas in m²
    """
    R = 6371000  # Earth radius in meters

    lats = data.lat.to_numpy()
    lons = data.lon.to_numpy()

    # Compute grid spacing (assuming uniform spacing)
    dlat = np.abs(np.diff(lats).mean())
    dlon = np.abs(np.diff(lons).mean())

    # Convert to radians
    dlat_rad = np.deg2rad(dlat)
    dlon_rad = np.deg2rad(dlon)
    lats_rad = np.deg2rad(lats)

    # Compute cell areas: A = R² × dlon × (sin(lat + dlat/2) - sin(lat - dlat/2))
    areas = np.zeros((len(lats), len(lons)))
    for i, lat_rad in enumerate(lats_rad):
        lat_north = lat_rad + dlat_rad / 2
        lat_south = lat_rad - dlat_rad / 2
        area = R**2 * dlon_rad * (np.sin(lat_north) - np.sin(lat_south))
        areas[i, :] = area

    return torch.from_numpy(areas)

def get_inference_steps(
    data_source: CanonicalSource,
    input_steps: int,
    output_steps: int,
):
    """
    Get the number of inference/rollout steps for the given time configuration.

    Args:
        data_source: The data source sliced to the inference time range
        input_steps: Raw timesteps consumed by each model call.
        output_steps: Future raw timesteps emitted by each model call.

    Returns:
        num_steps: Total number of rolled-out inferences which fit into the time range
    """
    available_targets = max(data_source.time.size - input_steps, 0)
    available_targets -= available_targets % output_steps
    # Aggregators record the initial input history followed by forecast targets.
    return input_steps + available_targets

def get_anomalies_vars(var_names: BoundaryVarNames) -> tuple[str, ...]:
    """Get the variables that need to be computed for anomalies."""
    return tuple([var for var in var_names if var.endswith("_anomalies")])

def compute_anomalies(
    data: xr.Dataset,
    means: xr.Dataset,
    stds: xr.Dataset,
    anomalies_vars: tuple[str, ...],
) -> tuple[xr.Dataset, xr.Dataset, xr.Dataset]:
    """Compute anomalies for the given variables."""
    for var in anomalies_vars:
        base_var = var.replace("_anomalies", "")
        if var not in data.variables and base_var in data.variables:
            logger.info(f"Computing anomalies for {base_var}")
            climatology = (
                data[base_var].groupby("time.dayofyear").mean("time").compute()
            )
            # Remove the seasonal cycle (climatology) from the detrended data
            day_of_year = data[base_var]["time"].dt.dayofyear
            data[var] = (
                data[base_var] - climatology.sel(dayofyear=day_of_year)
            ).compute()
            data = data.drop_vars(["dayofyear"])
            means[var] = data[var].mean().compute()
            stds[var] = data[var].std().compute()
    return data, means, stds

def with_level_index_vars(
    data: xr.Dataset,
    depth_levels: Sequence[float],
) -> xr.Dataset:
    """
    Ensure variable names use a depth level index, not depth level value.
    """
    data_copy = data.copy()

    for var in data.variables:
        # OM4 data format has variables in the form: var_lev_{depthlevel}
        # ex. so_lev_1040_0. We need to convert into var_{depthlevelidx}
        var_str = str(var)
        if "_lev_" in var_str:
            var_split = var_str.split("_lev_")
            var = var_split[0]
            lev_in_depth = float(var_split[1].replace("_", "."))
            lev_in_depth_idx = depth_levels.index(lev_in_depth)
            data_copy = data_copy.rename({var_str: f"{var}_{lev_in_depth_idx!s}"})

    return data_copy

def with_depth_value_vars(
    data: xr.Dataset,
    data_layout: DataLayout,
) -> xr.Dataset:
    """Inverse of `with_level_index_vars`: name 3D variables by depth value.

    Renames the depth-resolved prognostic variables (``<var>_<level_index>``)
    back to the OM4 ``<var>_lev_<depth>`` form (e.g. ``thetao_0`` ->
    ``thetao_lev_2_5``). Which variables to rename is read directly off
    ``data_layout.prognostic_var_names`` rather than inferred from the data, so
    per-level masks (``mask_<i>``) and level-free prognostics (e.g. ``zos``) are
    never mistaken for depth-resolved variables by name alone.
    """
    renames = {}
    for var_name in data_layout.prognostic_var_names:
        base, _, idx = var_name.rpartition("_")
        if not (base and idx.isdigit()):
            continue  # level-free prognostic variable, e.g. zos
        depth = data_layout.depth_levels[int(idx)]
        depth_str = str(depth).replace(".", "_")
        if var_name in data.variables:
            renames[var_name] = f"{base}_lev_{depth_str}"

    return data.rename(renames)

def with_lat_lon_coords(data: xr.Dataset) -> xr.Dataset:
    """Standardize dataset coordinates; prefer "lat"/"lon" over "y"/"x"."""
    data_copy = data.copy()
    if "lat" not in data_copy.dims:
        # Preserve the 2-D geographic coords under a non-colliding name, then
        # rename the x/y dims to the 1-D lat/lon we standardize on.
        preserve = {n: f"{n}_2d" for n in ("lat", "lon") if n in data_copy.coords}
        data_copy = data_copy.rename(preserve).rename({"x": "lon", "y": "lat"})

    return data_copy

def stack_levels(
    data: xr.Dataset,
    data_layout: DataLayout,
) -> xr.Dataset:
    """Reassemble a flattened OM4 dataset into analysis-ready, depth-stacked form.

    Inverts the preprocessing flattening so downstream analysis does not have to:
    per-level prognostic channels (``thetao_0`` ...) become ``thetao(lev, ...)``
    and the per-level ``mask_i`` become a single stacked ``wetmask``. Grid
    coordinates are preserved. This is the dataset-level counterpart to the eval
    writer's reassembly, for putting ground-truth inputs on the same footing as
    predictions.

    Implemented as the inverse of `with_level_index_vars` followed by the existing
    `compact_dataset` (stacks the ``_lev_`` form) and `unflatten_masks`.
    """
    data = with_depth_value_vars(data, data_layout)
    data = compact_dataset(data)
    if "mask_0" in data.variables:
        data = unflatten_masks(data, num_levels=len(data_layout.depth_levels))
    return data