Expose mesh layout selection in the create_graph_with_wmg CLI

[prajwal-tech07]

Following up on the Slack discussion with @leifdenby about extending the wmg–neural-lam bridge to support different mesh layouts.

Current state

The create_graph_with_wmg CLI (from #596) only exposes --archetype (keisler / graphcast / hierarchical). Each of those archetype functions in wmg hardcodes mesh_layout="rectilinear" internally and only varies the m2m_connectivity (flat / flat_multiscale / hierarchical). So even though wmg's create_all_graph_components takes a mesh_layout argument, there's no way to reach it from the neural-lam side today, and the CLI can only ever produce rectilinear graphs.

This issue is about exposing the mesh_layout axis (where the mesh nodes are placed) in the CLI, so users can build triangular and, later, prebuilt meshes through neural-lam.

Proposed approach

Instead of threading a mesh_layout argument through each archetype function, call wmg.create.create_all_graph_components directly from create_graph_from_datastore, passing mesh_layout and m2m_connectivity as separate arguments. This matches the two-step design the archetypes are built on (layout decides where nodes go, connectivity decides how they connect) and avoids confusing combinations like "keisler but triangular", where the archetype name no longer describes what was actually built.

In short:

• Add a --mesh_layout argument to the CLI.
• Still expose the connectivity choice (see open question below).
• Pass both into create_all_graph_components.

Driving the layout choices from wmg

To avoid hardcoding the list of valid layouts in neural-lam (which would drift as wmg adds new ones), @leifdenby suggested wmg expose the options as a public constant, e.g. wmg.create.MESH_LAYOUT_OPTIONS, just a list of strings. The CLI would use that to populate the --mesh_layout choices, so it expands automatically as wmg gains layouts, with no change needed on this side. This is consistent with how wmg already validates m2m_connectivity against a supported set internally, so it would also tidy up the current mesh_layout check (which is a hardcoded if == "rectilinear" right now). That part is a small wmg-side change and can be tracked separately.

Per-layout arguments

The layouts don't all take the same inputs:

• rectilinear and triangular are generated from grid spacing, so they only need mesh_node_distance (already handled by the CLI).
• prebuilt is different. The user has to supply the actual mesh node coordinates (for example an ICON or MPAS grid), which can't be derived from grid spacing. The CLI currently only reads grid coordinates from the datastore, so prebuilt would need a new input path such as --mesh_coords_path.

Because of that, I'd suggest doing triangular first (it's just a flag plus mesh_node_distance) and treating prebuilt as a follow-up once the input-source question is settled.

Validation

Each layout's output should still pass the graph storage validator (validate_graph from #323), so I'd extend the graph-creation tests to run the validator on each layout we expose. That keeps every layout verified end to end rather than just assumed to work.

Dependencies

• triangular depends on Add mesh_layout='triangular' for regular triangular mesh generation weather-model-graphs#92
• prebuilt depends on Add mesh_layout='prebuilt' for user-provided mesh graphs weather-model-graphs#91
• rectilinear already works, so the CLI plumbing itself can land independently of those.

Open questions

• Keep --archetype as a connectivity preset and add --mesh_layout next to it, or move fully to --mesh_layout + --m2m_connectivity?
• Where should MESH_LAYOUT_OPTIONS live in wmg, e.g. create/__init__.py?

[leifdenby]

Thanks for typing this up @prajwal-tech07!

Could you give a few examples of calling the CLI to create the 3 archetypes with the default (rectilinear) layout? And a fourth example where you change the mesh coordinate layout? E.g. examples something like:

uv run -m neural_lam.create_graph_with_wmg ...

[prajwal-tech07]

Sure! Here are some CLI usage examples.

Current behaviour - the three archetypes, all on the default rectilinear layout (this works today):

Keisler (flat single-scale)

uv run -m neural_lam.create_graph_with_wmg \
    --config_path path/to/config.yaml \
    --archetype keisler \
    --name keisler

GraphCast (flat multiscale)

uv run -m neural_lam.create_graph_with_wmg \
    --config_path path/to/config.yaml \
    --archetype graphcast \
    --name graphcast

Oskarsson hierarchical

uv run -m neural_lam.create_graph_with_wmg \
    --config_path path/to/config.yaml \
    --archetype hierarchical \
    --name hierarchical

Proposed addition (what this issue is about) - a --mesh_layout flag so you can keep the same connectivity but change where the mesh nodes are placed. For example, a keisler-style graph on a triangular mesh instead of rectilinear:

uv run -m neural_lam.create_graph_with_wmg \
    --config_path path/to/config.yaml \
    --archetype keisler \
    --mesh_layout triangular \
    --name keisler_triangular

Today the archetypes hardcode mesh_layout="rectilinear" internally, so --mesh_layout is the new part — it would default to rectilinear (current behaviour unchanged) and accept the other layouts as wmg adds them (driven by wmg.create.MESH_LAYOUT_OPTIONS, as discussed).

On prebuilt - agree with your Slack note, I'll leave it out of this CLI for now. It's enough of a special case (the user has to supply the actual mesh coordinates) that it's probably better as its own CLI rather than bending this one around it. So the layouts in scope here are the grid-derived ones (rectilinear, triangular).

[leifdenby]

This looks perfect! Let's go with this plan :)