Planar symmetries with staggered grids #99
Description
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!