Adding support for finite time differencing

[anthony-meza]

If we add time as a center coordinate and time_bounds as an outer coordinate to an xgcm.Grid call, we will be able to use difference to calculate monthly Eulerian time derivatives from snapshot diagnostics.

difference relies on grid.diff, and grid.diff can fail when the spatiotemporal axis being differenced is not contained in a single dask chunk. In practice, that means users may need to set chunks=-1 along the differencing axis (e.g. chunks={dim: -1} when calling grid.diff(..., axis=dim)).

Once time is added as a coordinate, MODEL.yaml graphs will be able to call difference along multiple spatiotemporal axes. This could unintentionally require the user to set chunks=-1 across many dimensions, which can accidentally pull the entire dataset into memory.

For these reasons, I propose the following features to be added to xbudget:

• (A) Add an optional allow_rechunk argument to collect_budgets that (if needed) will rechunk only the correct axis inside difference. This would avoid having to load the entire dataset into memory. If allow_rechunk=False, we can add a warning explaining that grid.diff needs the differencing axis to have chunks={dim: -1} along that axis.
• (B) Allow difference to be composed with other xbudget operations (I.e., difference(product(x, y))).
• (C) Add xgcm time coords to the setup of the MOM6 xgcm.Grid setup.

Currently typing this out on the flight to Glasgow lol, any thoughts/questions/concerns @hdrake ?

[hdrake]

While there are reasons we may want to diff in any dimension, they will generally be applied to different arrays, right? We don't need the -1 chunks to apply to every dimension for all variables in the dataset at once, right?

That said, I am concerned that users will lose track of what the chunking is for the different terms in the budget (and this gets worse in xwmb where we use a lot more grid operations).

Before discussing a way forward, I'd like to better understand how chunking interacts with grid.diff.

Would it be safe to do something like this?

grid.diff(da.chunk({"X":1, "Y"1, "Z", 1, "T":-1}), "T").chunk(original_chunks)

Maybe it would be safe but very inefficient? I like the idea that we undo all of the rechunking so that the end product is back in the original chunks that the user started with.

[anthony-meza]

"While there are reasons we may want to diff in any dimension, they will generally be applied to different arrays, right?"
Yes, generally diff will be applied to separate arrays. So it's a good idea to keep the -1 chunks to only the necessary dimensions.

I think restoring the original dataset's chunks is a good idea, and definitely safer. It's not clear to me yet whether it is more efficient. In my ECCO case, rechunking as a preprocessing step was unavoidable when applying difference (i.e., grid.diff) along a dimension that was already chunked along that same axis.

Here is some example code that I’ve confirmed works with the MOM6 example.

if k == "difference":
    if grid is not None:
        staggered_axes = {
            axn:c for axn,ax in grid.axes.items()
            for pos,c in ax.coords.items()
            if pos!="center"
        }
    v_term = [v_term for k_term,v_term in v.items() if k_term!="var"][0]
    if v_term not in ds:
        warnings.warn(f"Variable {v_term} is missing from the dataset `ds`, so it is being skipped. To suppress this warning, remove {v_term} from the `xbudget_dict`.")
        continue

    candidate_axes = [axn for (axn,c) in staggered_axes.items() if c in ds[v_term].dims]
    if len(candidate_axes) == 1:
        axis = candidate_axes[0]
    else:
        raise ValueError("Flux difference inconsistent with finite volume discretization.")

    if allow_rechunk: #NEW CODE
            try: #extract original chunks when possible
                #not using ds[v_term] since it may not have the non-staggered dimension chunks. 
                original_chunks = dict(ds.chunksizes)
            except Exception:
                warnings.warn("Dataset chunks are inconsistent; using unify_chunks()", UserWarning)
                original_chunks = dict(ds.unify_chunks().chunksizes)

            # Find the staggered dimension for the given axis in the DataArray
            axis_dim = [d for d in ds[v_term].dims if d in grid.axes[axis].coords.values()]
            if len(axis_dim) != 1:
                raise ValueError(f"Expected to find one dimension for axis '{axis}' in variable '{v_term}', but found {len(dims_for_axis)}: {dims_for_axis}")
            axis_dim = axis_dim[0]
        
            # Temporarily rechunk to put the difference dim in a single chunk, all other chunks are auto. 
            temporary_chunks = {axis_dim: -1, **{d: "auto" for d in ds[v_term].dims if d != axis_dim}}
            var = grid.diff(ds[v_term].chunk(temporary_chunks).fillna(0.0), axis=axis)
            # Attempt original chunking for preserved dimensions
            var = var.chunk({d: original_chunks.get(d, var.chunksizes[d]) for d in var.dims})
    else:  #OLD CODE
            var = grid.diff(ds[v_term].fillna(0.), axis)

    var_name = f"{namepath}_difference"
    var = var.rename(var_name)
    var_provenance = v_term
    var.attrs["provenance"] = var_provenance
    ds[var_name] = var
    if var_pref is None:
        var_pref = var.copy()
    else:
        raise ValueError("Input `ds` must be `xgcm.Grid` instance if using `difference` operations.")

[hdrake]

This seems sensible to me! No rush, but would be great if you could start a PR for this feature and include some basic tests.

I'll want to try it out myself and see how this affects dask graphs and performance before we merge it.

[anthony-meza]

See #24 for a pull request address #21 (A)