Icon grids support

[liuquan18]

ICON grids are not so different from HEALpix grids in terms of multi-scale decomposition.

1. ICON grids, each parent triangle has four child triangles, similar to each HEALpix parent pixel has four child pixels; In ICON grids, n_cells(R2Bk) = 20 · 4^(k+1)/4 = 80 · 4^k (equivalently 20·2²·4^k) — the same "fixed base-tile count, then exact ×4 per level" shape as HEALPix's 12·4^z, just with 20 base icosahedron faces instead of 12 base pixels.

2. ICON grids are stored in a "nested" way, these 4 children occupy contiguous indices, always ordered [center, corner1, corner2, corner3]. The cell centre that corresponds to the corner1 may different in different cells. but the order is not really important for coarsening and refining.

3. further validation shall that the result should be mathematically correct.

Therefore, three new functions are added in the grid_utils.py : icon_neighbor_cell_index_to_adjc, which corresponds to healpix_get_adjacent_cell_indices; icon_grid_to_mgrid which corresponds to healpix_grid_to_mgrid, and a new function to make sure the nested layout validate_nested_ordering.

The following scripts can be used for testing, and the test_data (R2B4 fine grids, R2B3 coarse grids, and an example data with variable ['u_10m']) can be found here /work/mh0033/m300883/data_share/test_data.

import math
import numpy as np
import torch
import torch.nn as nn
import xarray as xr

from fieldspacenn.src.modules.field_space.field_space_attention import FieldSpaceAttentionConfig
from fieldspacenn.src.modules.field_space.field_space_base import ConservativeLayerConfig
from fieldspacenn.src.models.mg_transformer.mg_transformer import MG_Transformer
from fieldspacenn.src.modules.grids.grid_utils import icon_neighbor_cell_index_to_adjc, icon_grid_to_mgrid, validate_nested_ordering

torch.manual_seed(0)



# ---------------------------------------------------------------------------
# 2. Build the grid hierarchy + model (mirrors tutorial Section 7.2.2)
# ---------------------------------------------------------------------------

DATA_DIR = "test_data"
COARSE_GRID = f"{DATA_DIR}/Earth_IcosS_0320km.nc"   # R2B3, 5120 cells -> internal level 0
FINE_GRID = f"{DATA_DIR}/Earth_IcosS_0160km.nc"     # R2B4, 20480 cells -> internal level 1
DATA_FILE = f"{DATA_DIR}/nwp_r2b4_atm_2d_ml_19791201T000000Z.nc"

assert validate_nested_ordering(COARSE_GRID, FINE_GRID)

mgrids = icon_grid_to_mgrid({0: COARSE_GRID, 1: FINE_GRID})
n_coarse = mgrids[0]["coords"].shape[0]
n_fine = mgrids[1]["coords"].shape[0]
print(f"coarse (level 0, R2B3) cells: {n_coarse}   fine (level 1, R2B4) cells: {n_fine}")

block_configs = {
    "0": FieldSpaceAttentionConfig(
        token_zoom=0, q_zooms=[0, 1], kv_zooms=[0, 1],
        att_dim=32, n_head_channels=16,
    ),
    "1": FieldSpaceAttentionConfig(
        token_zoom=0, q_zooms=[0, 1], kv_zooms=[0, 1],
        att_dim=32, n_head_channels=16,
    ),
    "conservative": ConservativeLayerConfig(),
}

model = MG_Transformer(
    mgrids=mgrids,
    block_configs=block_configs,
    in_zooms=[0, 1],
    in_features=1,
    n_groups_variables=[1],
    n_head_channels=16,
)
n_params = sum(p.numel() for p in model.parameters())
print(f"model built: {n_params} parameters")

# ---------------------------------------------------------------------------
# 3. Real data: u_10m on the R2B4 grid, 24 hourly snapshots -> 24 training samples
# ---------------------------------------------------------------------------

ds = xr.open_dataset(DATA_FILE)
u_fine_all = torch.from_numpy(ds["u_10m"].isel(height=0).values).float()   # (time=24, n_fine)
print("u_10m raw shape:", u_fine_all.shape)

# low-res input: mean-pool the real fine field down to the coarse grid (Section 2 of
# docs/icon_native_grid_reshape_feasibility.md justifies using .mean(-1), not [...,::4])
u_coarse_all = u_fine_all.view(u_fine_all.shape[0], -1, 4).mean(dim=-1)   # (time=24, n_coarse)

# simple standardization (fit on this sample; fine for a demo)
mean, std = u_fine_all.mean(), u_fine_all.std()
u_fine_n = (u_fine_all - mean) / std
u_coarse_n = (u_coarse_all - mean) / std

def to_zooms_batch(coarse_batch, fine_batch):
    """(B, n) tensors -> the {zoom: (b,v,t,n,d,f)} dict layout MG_Transformer expects."""
    B = coarse_batch.shape[0]
    return {
        0: coarse_batch.view(B, 1, 1, -1, 1, 1),
        1: fine_batch.view(B, 1, 1, -1, 1, 1),
    }

# ---------------------------------------------------------------------------
# 4. One training epoch over the 24 samples
# ---------------------------------------------------------------------------

n_samples = u_fine_n.shape[0]
batch_size = 4
optimizer = torch.optim.Adam(model.parameters(), lr=2e-4)
loss_fn = nn.MSELoss()

perm = torch.randperm(n_samples)
sample_configs = {z: {"n_past_ts": 1, "n_future_ts": 1, "zoom_patch_sample": -1, "mask_n_last_ts": 1} for z in (0, 1)}

epoch_losses = []
model.train()
for start in range(0, n_samples, batch_size):
    idx = perm[start:start + batch_size]
    x_coarse = u_coarse_n[idx]
    y_fine = u_fine_n[idx]

    x_zooms_groups = [to_zooms_batch(x_coarse, torch.zeros_like(y_fine))]  # fine zoom unknown -> zero-init
    out = model(x_zooms_groups=x_zooms_groups, emb_groups=[{}], sample_configs=sample_configs, out_zoom=1)
    pred_fine = out[0][1].view(y_fine.shape)

    loss = loss_fn(pred_fine, y_fine)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    epoch_losses.append(loss.item())
    print(f"batch [{start:2d}:{start+len(idx):2d}]  MSE={loss.item():.4f}")

print(f"\nepoch mean MSE: {np.mean(epoch_losses):.4f}")

# sanity: naive baseline (just re-upsample the coarse input, no learning at all)
naive_pred = u_coarse_n.repeat_interleave(4, dim=-1)
naive_mse = loss_fn(naive_pred, u_fine_n).item()
print(f"\nnaive upsample-only baseline MSE: {naive_mse:.4f}  (model's epoch-1 MSE was ~{np.mean(epoch_losses):.4f}, expected to start near this)")

print("\n--- extended run: 40 more epochs, to confirm the model actually learns beyond the naive baseline ---")
for epoch in range(40):
    perm = torch.randperm(n_samples)
    losses = []
    for start in range(0, n_samples, batch_size):
        idx = perm[start:start + batch_size]
        x_coarse = u_coarse_n[idx]
        y_fine = u_fine_n[idx]
        x_zooms_groups = [to_zooms_batch(x_coarse, torch.zeros_like(y_fine))]
        out = model(x_zooms_groups=x_zooms_groups, emb_groups=[{}], sample_configs=sample_configs, out_zoom=1)
        pred_fine = out[0][1].view(y_fine.shape)
        loss = loss_fn(pred_fine, y_fine)
        optimizer.zero_grad(); loss.backward(); optimizer.step()
        losses.append(loss.item())
    if epoch % 5 == 0 or epoch == 39:
        print(f"epoch {epoch+2:3d}: mean MSE = {np.mean(losses):.4f}")

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 481b7c86e8

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

Resolve ICON coordinates consistently in validator

When grid files use clon/clat (which icon_grid_to_mgrid explicitly accepts), this function raises KeyError instead of returning an ordering verdict because it only reads lon_cell_centre/lat_cell_centre. That prevents the advertised preflight check for an otherwise supported ICON grid pair; resolve the coordinate variable names using the same fallback logic as the grid builder.

Useful? React with 👍 / 👎.

[liuquan18]

maybe something like

         lon_name = "clon" if "clon" in ds else "lon_cell_centre"
         lat_name = "clat" if "clat" in ds else "lat_cell_centre"

[chatgpt-codex-connector]

Validate all children before approving nested ordering

For a candidate grid where entry 4*i is near coarse cell i but entries 4*i+1:4*i+4 have been shuffled among parents, this calculation still returns True. The purported guard would then approve reshape-based coarsening that mixes unrelated cells, because it checks only the first child in each four-cell block; verify all four children (and their group membership) before reporting success.

Useful? React with 👍 / 👎.