Support vector-valued UxDataArrays in vector calculus operations

Attempting to solve #1878 by implementing vector calculus methods into UxDataset directly from their current form raises lots of questions:

1. gradient() outputs a UxDataset already, with "zonal_gradient" and "meridional_gradient" data vars. So:

• (1a) What should ds.gradient() look like?
• (1b) If ds has data vars "T", "P", should the result have twice as many vars, something like "T_zonal_gradient", "T_meridional_gradient", "P_zonal_gradient", "P_meridional_gradient"?
• (1c) Or, should the output be two UxDataArrays, one with zonal gradients and the other with meridional gradients?

2. curl(), divergence(), and scalardotgradient() take the zonal and meridional components as separate input arrays. So:

• (2a) What would ds.curl() mean?
• (2b) Or, should the API look like u_ds.curl(v_ds) instead, where u_ds is a dataset of zonal components and v_ds is a dataset of meridional components?
• (2c) Should uxarray allow u_ds.curl(v) where v is a single UxDataArray, and/or u.curl(v_ds)?
• (2d) How can one unambiguously provide the scalar field q into a call to ds.scalardotgradient()?

Instead of answering these questions directly, this issue focuses on proposing a new option for existing vector calculus methods that would remove the need to ever answer most/all of these design questions in the first place.

Proposal

My proposed solution is to add an option throughout the UxDataArray vector calculus methods which allows utilizing arrays with zonal and meridional components are stacked together into a single array with a vector dimension, instead of forcing them to be treated separately. For example, one could provide a wind array stacked along a "component" dimension, instead of two separate arrays for wind_zonal and wind_meridional. This option could be implemented in a way which is fully backwards compatible; here is a rough sketch:

def gradient(self, scale_by_radius=True, *, vector_dim=None):
    """returns a UxDataset with "zonal_gradient" and "meridional_gradient" data vars.
    If vector_dim is provided, instead return a UxDataArray,
    with zonal and meridional results stacked along this dim."""
    # ...

def curl(self, other=None, scale_by_radius=True, *, vector_dim="component"):
    """returns a UxDataArray telling curl(self, other).
    If other is None, self must contain vector_dim with "zonal" and "meridional" coords;
    zonal and meridional components will be inferred as self.sel({vector_dim: "zonal"})
    and self.sel({vector_dim: "meridional"})."""
    # ...

def divergence(self, other=None, scale_by_radius=True, *, vector_dim="component"):
    # similar to curl()

For scalardotgradient, I would maybe not add an option like this at all, due to a bit of ambiguity / reduced benefits due to the additional scalar field input. Also, when utilizing the "stack along vector dimension" option, the arithmetic of the scalardotgradient functionality becomes an extremely simple one-liner anyway: wind.sel(component="zonal").scalardotgradient(wind.sel(component="meridional"), q) can be expressed simply as (wind * q.gradient()).sum("component").

This option would make it easy to add a UxDataset implementation in #1873 because it changes the inputs/outputs to always be a single UxDataArray, like "UxDataArray --> UxDataArray" for gradient, curl, and divergence. (As opposed to their current inputs/outputs, which look like "UxDataArray --> UxDataset" for gradient, and "(UxDataArray, UxDataArray) --> UxDataArray" for curl and divergence.)

Additional Benefits

I noticed these points of friction a while ago but hadn't created an issue report because they aren't necessarily convincing enough on their own to fully warrant maintaining a new feature / alternative syntax. But, they would come along with this new feature for free:

• Users who use vector_dim no longer need to remember the special names "zonal_gradient" and "meridional_gradient" for the data vars output by gradient()
• Users who use vector_dim no longer have any chance to accidentally swap the argument order for something like A.curl(B).
• It becomes easier to chain operators together, such as in the scalardotgradient example above. Another example could be a drop-in replacement for the Laplacian like T.gradient(vector_dim="component").divergence(). The equivalent syntax users currently have available to them looks like: g = T.gradient(); g["zonal_gradient"].divergence(g["meridional_gradient"]).

Misc. notes

There is some xarray precedent for having a dimension be "the vector dimension" such as in xr.cross(a, b, dim=...) which computes cross products along a named component dimension which the inputs must be stacked along (e.g. a is really a_x, a_y, a_z stacked along dim).

Maybe some day the default for gradient() might shift to be to output a single UxDataArray along the vector_dim, and the default vector_dim could be component, but that would require a deprecation cycle, might not be desired, and is out of scope for this issue.

AI disclosure: The proposal here looks very close to my initial thoughts on the topic, but I did work through ideas extensively with Claude (Opus 5.5), including various possible alternatives and figuring out which pieces to highlight here, with the goal of narrowing in on a proposal that is most likely to be appealing to the uxarray dev team, based on my understanding of team preferences. I wrote up this issue report myself; it's long because there is a lot of content, not due to extra verbose AI rambling.