Planar symmetries with staggered grids

[dgreenberg]

Hi, and first of all thanks for creating this fantastic package!

We're interested in using escnn to build ML-based solvers for fluid dynamics. In our applications, the inputs typically consist of vector fields with marker-and-cell staggering (Arakwa c-grids, see image below). This means that if pressure/height variables are at integer locations, then horizontal velocities u are offset by half a grid cell to the horizontally, and vertical velocities v by half a grid cell vertically. Examples include 2D Navier Stokes and the shallow water equations.

We've verified mathematically and numerically that the reference solvers we use to generate training data are equivariant to machine precision. However, it's not so clear how we can apply escnn for our staggered grid. For example, if we associate the (u, v) pair to the right and above each pressure point as a single vector in our vector field, then after counterclockwise rotation by 90 degrees the u vector becomes a new v vector and the v vector changes sign and becomes a u vector. But they would now no longer be associated with each other, but with other u and v values. This means that the group no longer operates on the fiber bundle at each point in field, but rather exchanges values between nearby finder bundles.

If I'm not mistaken, this means that escnn can't be used out of the box with this grid staggering. But perhaps I've missed something? I'd be grateful for any guidance on this point!

[maxxxzdn]

It seems to be a rather non-trivial setting, and I don't think escnn has an out-of-box solution for that.

That being said, perhaps you can reformulate the task a little bit to make it compatible:

1. define a feature field $u$ discretized on grid $x_1$, and a feature field $v$ discretized on grid $x_2$.

u_gt = field_type(u_data, x_1)
v_gt = field_type(v_data, x_2)

2. apply two separate encoders

z_u = self.encoder_u(u_gt)
z_v = self.encoder_v(v_gt)

3a. use interpolation for exchanging information between fibers on $x_1$ and $x_2$ OR
3b. interpolate the latent fibers to a common grid $x_3$, continue the computation on it.

This formulations should be equivariant, though I am not sure how convenient it is for your task.