Simple "Difference" Operator

[jejjohnson]

Is there a native way to do a finite difference operator on a multidimensional field, i.e. $\partial_x \vec{\boldsymbol{u}}, \partial_y \vec{\boldsymbol{u}}, \partial_z \vec{\boldsymbol{u}}$, ....

Using the current API, I don't see a way to specify the scheme. It could a combination of the stagger option but I could not find a way to do it. There is a gradient operation but many times we just need a simple difference operator where we can choose which axis.

Note: This might be a problem due to the fact that we use convolutional operators for the FD scheme whereas normally I think of slicing.

---

Demo

I have a demo colab notebook to showcase what I mean. The equation of motion is a 1D problem but if it were 2D then this API would not work.

---

Proposed Solution

No specific solution but it might be helpful to have a simple API for this as in many models like the Shallow water and Quasi-geostrophic models, we need this because we have a lot of advection terms.

# current API
u_grad = gradient(u)
du_dx = u_grad.replace_params(u_grad.on_grid()[0])

# preferred API
du_dx = difference(u, axis=0)

Another solution is just to write a custom operator for the difference scheme. The colab notebook that I linked before has an example of this. This is also related to issue #127 and #125.

[astanziola]

I might have misunderstood, but basically what you are looking for is a jacobian operator?

That's in the following sense (I'll write in in 2d for simplicity).
Given a field $u(r) = (u_x(r), u_y(r))$ with $r \in \mathbb R^2$, you are looking for

$$ \nabla u = \begin{pmatrix} \frac{\partial u_x}{\partial x} & \frac{\partial u_x}{\partial y} \\ \frac{\partial u_y}{\partial x} & \frac{\partial u_y}{\partial y} \end{pmatrix} $$

If that's the case, we could add a jacobian operator, which is probably useful in general. Here's how it looks for finite differences: https://colab.research.google.com/notebooks/intro.ipynb