`create_graph` always generates square mesh grids, ignoring data grid aspect ratio

[osten-antonio]

Problem

In line 244-259, the mesh graph is always square regardless of the data grid dimensions:

# line 243
nx = 3  # number of children =nx**2
nlev = int(np.log(max(xy.shape[:2])) / np.log(nx)) # Uses largest dimension
nleaf = nx**nlev  

# line 258
n = int(nleaf / (nx**lev))
g = mk_2d_graph(xy, n, n) # Same n for both x and y

For a data grid of shape (Nx, Ny), the mesh is always n x n, where n is derived from max(Nx, Ny). In addition, the dm calculation only measures the x-direction, assuming that the x direction is the largest:

dm = np.sqrt(
    np.sum((vm.data("pos")[(0, 1, 0)] - vm.data("pos")[(0, 0, 0)]) ** 2)
)

Since the mesh may not be uniform, this wont correctly capture the y-direction relationship.

Why it matters

Using the largest dimension means that for a case where the data grid has a non-square aspect ratio will produce multiple redundant node, wasting computation. This also comes with the added side effect of having uneven spacing, meaning that the weather moving diagonally will look different depending on the direction relative to the mesh, because it treats x and y as different resolution.

For example, for a 200x30 domain (6.7:1 aspect ratio), the original code produces a 27x27 mesh where y-nodes would be 6.7x more densely packed than x-nodes.

How to reproduce/visualize

The following is the script that is used to visualize the problem

import torch
import tempfile
import os
import numpy as np
from neural_lam.create_graph import create_graph

Nx, Ny = 200, 30
x = np.linspace(0, Nx, Nx)
y = np.linspace(0, Ny, Ny)
xx, yy = np.meshgrid(x, y, indexing='ij')
xy = np.stack([xx, yy], axis=-1)

print(f"Data grid: {Nx}x{Ny} (aspect ratio {Nx/Ny:.1f}:1)")

with tempfile.TemporaryDirectory() as tmp:
    create_graph(tmp, xy)

    mesh = torch.load(os.path.join(tmp, 'mesh_features.pt'))
    pos  = mesh[0]

    x_unique = len(torch.unique(pos[:, 0]))
    y_unique = len(torch.unique(pos[:, 1]))

    dx_phys = Nx / x_unique
    dy_phys = Ny / y_unique

    print(f"Mesh shape: {x_unique}x{y_unique} = {pos.shape[0]} nodes")
    print(f"x spacing: {dx_phys:.1f} units/node")
    print(f"y spacing: {dy_phys:.1f} units/node")
    print(f"Spacing ratio: {dx_phys/dy_phys:.1f}:1  (ideal: 1.0:1)")
    print(f"Redundant y nodes: ~{y_unique - round(y_unique * Ny/Nx)}"
          f" of {y_unique} rows cover the same grid points")

The output given for 200x30 is:

Mesh shape: 27x27 = 729 nodes
x spacing: 7.4 units/node
y spacing: 1.1 units/node
Spacing ratio: 6.7:1
Redundant y nodes: ~23 of 27 rows cover the same grid points 

Proposed implementation

At each mesh level, both n_x and n_y should be scaled independently to match the physical aspect ratio of the domain, rather than using the same value n for both dimensions. The dm connection radius should also be updated to correctly reflect the non-square mesh cell geometry to ensure full grid coverage.

Related

Issue #4

[osten-antonio]

Hi @sadamov sorry for jumping straight to the proposed implementation, could I be assigned to this issue? Thanks!

[sadamov]

#299 #373 seem to already tackle this issue, or am I missing somethig?

[osten-antonio]

#299 #373 seem to already tackle this issue, or am I missing somethig?

Ah yea, I addressed it at #373

[sadamov]

#373 should have closed this automatically