Improve `XarrayLike` types to consolidate `xr.DataArray` and `xr.Dataset`

[nikeethr]

Is your feature request related to a problem? Please describe.

Having to deal with both xr.DataArray AND xr.Dataset leads to branching logic and additional tests all over the code.

Describe the solution you'd like

XarrayLike is a type that is accepted as input to most functions. I'd like to propose XarrayLike extended to unify the types, rather than just a static type check utilty. In essence it should act as a transformer on runtime to map the types to a single type. This way any code logic dealing with XarrayLike can operate under a single-type assumption.

I think the resultant unified datatype should be xr.Dataset because:

• xr.DataArray is, I believe, fully representable within a dataset. i.e. a dataset with one array. ¹
• xr.Dataset is already one half of the representable types of XarrayLike
• I'm ignoring pandas for now as out of scope...

risks:

• refactoring required and code cleaning
• need to make sure functionality is not affected: luckily we have good coverage

benefits:

• less code bloat.
• any common utility function only needs to deal with one type.
• less adhoc branching logic - each branch logic has added chance of error.
• tests can focus more on internal logic rather than repeatedly testing for handling the two types.

Describe alternatives you've considered

• Standardize to xr.DataArray instead ¹
• Honestly, if a xr.DataArray is fully representable by a xr.Dataset, then this is a hard solution to beat. Any alternative would have to be some other data structure that can interface between the two more smoothly than xr.Dataset inherently can.

Footnotes

1. The reverse is not true (i.e. xr.Dataset is not fully representable as a xr.DataArray) - and we will have to invent our own collection to standardise a dataset to a data array. ↩ ↩²

[nikeethr]

A couple of considerations:

• A reasonable number of xarray methods that I've randomly look at in the past that deal with xr.DataArray, under the hood dispatch to its xr.Dataset equivilent. This is either:
  • due to legacy OR
  • the functions natively being implemented for datasets because any indexing logic is more efficient especially if they get shared between data variables.
• If we go with the most basic implementation it would be:
  • lift data arrays to a datasets, can be done by a common utility. ¹
  • remove unnecessary logic that deals with data arrays and simplify any branched operations.
  • for cases only provided by data arrays we apply to the internal arrays one at a time (scores would have done this or similar for datasets already, so we just keep that).
  • simplify tests
• The more complex implementation would require a functor.
  • I won't go through the details, but basically we would now need XarrayLike to act as a container that also traces the original type (e.g EnumXarrayType.DATASET, EnumXarrayType.DATAARRAY). As well as lift any xarray functions to act on XarrayLike (e.g. via a wrapper).

The simple approach is probably the easiest to refactor into and test initially, since most functions should have already maintained compatibility to both datasets and dataarray. The only caveat is that the return type will always be a dataset even if the input is a data array, but maybe this is a minor concern.

The complex approach will guarantee the output type matches the input, but increases dev effort - mainly because instead of doing things like ds1 - ds2, it'd need to be routed through an applicator like XarrayLike.apply(xr.ufunc.subtract, ds1, ds2), which would lift the result back into an XarrayLike. Or alternatively XarrayLike(xr_data1.inner() - xr_data2_inner()). Neither of these are particularly pleasing or readable. One could think of subclassing from xr.Dataset but this will not guarantee preservation of the "original type" Enum when operations are performed on it, nor would xarray respect the structure of XarrayLike if its constructing its own xr.Dataset to return - which isn't subclassed. There maybe other ways to polymorph the design so that something can run after any method call associated to an object, but I'm not familiar enough with python or xarray to comment about that.

I'd prefer the simple approach if always returning a dataset is a non-issue. The other approach needs some functional programming knowledge and its parallels that can apply to python objects, in order to implement properly.

Footnotes

1. In fact if these common utilities (like gather_dimensions) are becoming common place that they are repeated in each function, instead of specifying them in each function, there can be a common dispatcher e.g. GenericScore which is can contain a method like lift_xr_data_to_xr_datasets(), that lift any data arrays, checks() that runs common checks, and score() where the actual score is implemented and calculated, and run() which does everything i.e. consolidates types, checks weights, gather dimensions and runs the score. There can be exceptions to the rule, but for simple enough scores, inheriting from such a structure will avoid repetition. ↩

[nikeethr]

Design

                        legacy behaviour    │         new behaviour
 ◄──────────────────────────────────────────┴─────────────────────────────────────────►
                                            │
  (legacy behaviour) result           input │ 
                       ▲                │   │
                       │                │   │
                       │                │   ┌────────────────────────────────────┐
                       │                │   │  internal structure                │
                       │                │   │                                    │
                       │                ▼   │  l_fn           r_da               │
                       │    ┌──────◄────┼───┼────►── X_AR ◄────────────┐         │
                       │    │ is X_AR?  │   │          │               │         │
               ┌── UTILS ◄──┤           │   │          │               │         │
               │            │           │   │     l_ar │               │         │
               │            │           ▼   │          ▼    l_ds       │         │
               │            └──────◄────┴───┼────►── X_DS ─────────► L_DS ──┐    │
               │              is X_DS?      │   l_fn   ▲               │    │    │
               │                            │          └───────────────┘    │    │
               │                            │                 r_ds          │    │
               │                            │                               │    │
               └─────────────►──────────────┼──────────► L_UTILS ◄──────────┘    │
                            l_fn            │                │                   │
                        (compatiblity)      │                │                   │
                                            └────────────────┼───────────────────┘
                                                             │
                                                             ▼
                                                           result (new behaviour)

   ┌─────────────────────────────────────────────────────────────────────────────────────────────────┐
   │ OBJECTS                                                                                         │
   │                                                                                                 │
   │ X_AR = xarray dataarray                                                                         │
   │ X_DS = xarray dataset                                                                           │
   │ L_DS = lifted dataset, that contains the original X_DS, or lifts X_AR -> X_DS while retaining   │
   │        any information required to undo this.                                                   │
   │ UTIL = "utility" function, but really any method that needs to distinguish between a dataarray  │
   │        or a dataset.                                                                            │
   │ L_UTIL = lifted version of the utility function that automatically handles lifted datasets      │
   └─────────────────────────────────────────────────────────────────────────────────────────────────┘

   ┌─────────────────────────────────────────────────────────────────────────────────────────────────┐
   │ MAPPINGS                                                                                        │
   │                                                                                                 │
   │ l_ar = lifts a data array to a dataset - this is essentially xarray's to_dataset                │
   │ l_ds = lifts a dataset to a "lifted" dataset (L_DS) which introduced by scores                  │
   │ r_da = reverts a lifted dataset  back to a xarray dataarray (L_DS -> X_AR)                      │
   │ r_ds = reverts a lifted dataset back to a xarray dataset (L_DS -> X_DS)                         │
   │ l_fn = compatibility helper to allow legacy functions to use L_DS instead of XarrayLike         │
   └─────────────────────────────────────────────────────────────────────────────────────────────────┘

---

Examples

1. Current pattern

    CURRENT PATTERN
    ---------------
                                                                     ─────────────┐
                                        ┌────────┬──── INTERNAL LOGIC (dataarray) │
                                        │        │                                │
                                        │        └──── INTERNAL LOGIC (dataset)   │
                                        │                                         │
                                        │                                         │
                                        │  ┌─────┬──── INTERNAL LOGIC (dataarray) │
                                        │  │     │                                │
                              ┌─ UTIL1 ─┘  │     └──── INTERNAL LOGIC (dataset)   │
                              │            │                                      │
                              ├─ UTIL2  ───┘                                      │
        Input ──────► SCORE ──┤               ──────────────┐                     │
                              ├─ INTERNAL LOGIC (dataarray) │                     │
                              │                             │◄────────────────────┘
                              ├─ INTERNAL LOGIC (dataset)   │
       Result ◄───────────────┴──────────◄──────────────────┘
                                                           

• logic needs to be branched to work for both dataset (X_DS) and (X_DA) EVERYWHERE that supports
  XarrayLike IF we want to maintain compatiblity
• This is especially troublesome to maintain if we allow multiple mixed input types

(psuedo)code example

# CURRENT PATTERN
#
# - input(s) & output is XarrayLike (dataset or dataarray)
#
# - EVERY function in the execution stack that either supports an input or output XarrayLike MUST
#   perform some form of branching logic IF datasets and dataarrays have divergent logic
#
# - This is because XarrayLike is a product type, and every combination between dataset and
#   dataarray API needs to be deterministic and mapped one to one and onto for reproducible
#   operation (this may not be the case - so its more realistic that some scores support one but not
#   the other)
#
# - APIs of datasets and dataarrays may diverge at any point, which makes the assumption that
#   they will have the same API compatiblity a risky design decision


def util1(x: XarrayLike, ...) -> XarrayLike:  # may return other types as well or just raise asserts
    assert_is_xarray(x)
    if isinstance(x, xr.Dataset):
         ...
    else:
         ...

def util2(x: XarrayLike, ...) -> XarrayLike:  # may return other types as well or just raise asserts
    assert_is_xarray(x)
    if isinstance(x, xr.Dataset):
         ...
    else:
         ...


def score(x: XarrayLike, ...) -> XarrayLike:
    assert_is_xarray(x)

    if isinstance(x, xr.Dataset):
        y = util1(x)
        assert_dataset(y)  # Most scores currently DO NOT currently do this
                           # the score itself does have no guarentee that the utility
                           # returns the same type it requires, and vice versa
        z = util2(y)
        assert_dataset(z)  # This may see repetitive - that's because it is and
                           # when dealing with product types it NEEDS to be done
                           # if we want to avoid nasty errors with mixed types.
    else:
         y = ...
         assert...
         z = ...
         assert ...
         ... same as above

    return z

# Important: If this branching is not consistent, then we have non-deterministic behaviour, whereby
# compatibility between dataarrays and datasets is unknowable at some point in the stack.

---

2. Proposed pattern - compatibility with legacy scores

    PROPOSED PATTERN (for legacy code - until refactor)
    ----------------
                                  l_fn                                l_fn
        Input ──────────► SCORE ───────── L_SCORE ────────┬─ UTIL1 ─────────────┐
                          │ ▲            (dataset)        │           l_fn      │
                          │ │                             ├─ UTIL2 ─────────────┤
                          │ │                             │                     │
       Result ◄───────────┘ │                             ├─ INTERNAL LOGIC     │
                            │                             │  (dataset only)     │
                            │                             │                     │
                            │             r_ds/r_da       ▼                     │
                            └──────────────────────────────────── ◄─────────────┘

• Here a compatibility decorator can be used for some of the utility functions before fully migrating to the full refactored solution.

(psuedo)code example

# COMPATIBILTY PATTERN ("bridging" solution)
#
# - input(s) & output is still XarrayLike (dataset or dataarray) - public API has not changed
#
# - no change in any logic in scores and internals
#
# - wrapper automatically lifts any dataset/dataarray inputs into a "LiftedDataset"
#
# - now the legacy functions can work with "LiftedDataset" objects automatically.
#
# - this guarentees automatic safety checks performed by lifting.
#
# - the wrapper is intended as a transitional helper prior to a full refactor.


@lift_to_ds
def util1(x: XarrayLike, ...) -> XarrayLike:
	... # no change to implementation

@lift_to_ds
def util2(x: XarrayLike, ...) -> XarrayLike:
	... # no change to implementation

@lift_to_ds
def score_ds(x: XarrayLike, ...) -> XarrayLike:
	... # move original score implementation, ignore anything to do with dataarrays or porting them
		# to work for datasets

def score(x: XarrayLike, ...) -> XarrayLike:
	return score_lifted(x, ...)  # dispatch to lifted score instead

---

3. Proposed pattern - new scores/refactoring old scores

Upon refactor, scores and utility should natively support lifted data frames. The main API does not change. However, scores offload/dispatch their logic to a "lifted" version that deals with LiftedDatasets. This way only logic to do with dataset needs to be maintained this:

• improves compatibility between dataarray and datasets
• reduces risk from requiring to maintain both
• reduces bloat, and likelihood of human error

    PROPOSED PATTERN (for new code)
    ----------------
                                          (internal lifting, instead of wrapper)
        Input ──────────► SCORE ───────── L_SCORE ────────┬─ L_UTIL1 
                          │ ▲            (dataset)        │           
                          │ │                             ├─ L_UTIL2
                          │ │                             │                
       Result ◄───────────┘ │                             ├─ INTERNAL LOGIC
                            │                             │  (dataset only)
                            │                             │                
                            │             r_ds/r_da       │                
                            └─────────────────────────────┘

(psuedo)code example

# PROPOSED PATTERN (desired solution) 
#
# - input(s) & output is still XarrayLike (dataset or dataarray) - public API has not changed.
#
# - underlying utilities need to replace XarrayLike with LiftedDataset.
#
# - logic in any score should be moved to an internal function that only works with LiftedDataset.
#
# - the main score API is responsible for converting any inputs to LiftedDataset and only using
#   utility functions compatible with LiftedDataset(s)
#
# - any branching logic between dataarrays and datasets can now be removed, as it is automatically handled

def util1(x: LiftedDataset, ...) -> LiftedDataset:
	assert_lifteddataset(x)  # instead of assert_xarraylike
	x_ds = x.ds              # retrieve reference to lifted dataset

	...  # directly implement logic for datasets;
	     # compatibility with dataarray-only API is done by iterating `x_ds.data_vars`

def util2(x: LiftedDataset, ...) -> LiftedDataset:
	assert_lifteddataset(x)
	x_ds = x.ds
	...

def score(x: XarrayLike, ...) -> XarrayLike:
	assert_xarraylike(x)	

	# no branching logic required...

	lds_x = LiftedDataset(x)
	lds_y = util1(lds_x)
	lds_z = util2(lds_y)
	return lds_z.raw()  # destroy outer wrapper and return internal `XarrayLike` structure
                            # - returns dataarray if originally dataarray, else dataset
                            # - this is automatically resolved with call to `.raw()`

Alternatives

Utility function to do branching logic on each function

• This only solves half the problem, each function would still need to know when to convert something to a dataset, and maintain the type of the original input. It will require several utility functions to perform and undo these steps.
• A lifted data structure abstracts this, by storing the underlying input data and its original state throughout its lifetime.

Do branching per function on a case by case basis

• This is the current design - case by case design is not pythonic - there should be one agreed upon way of doing it.
• Currently the branching logic itself is inconsistent for different functions, e.g. sometimes it is at a outer level, sometimes it is per function. Even if it is consistent, maintaining consistency as the codebase grows is tricky

Just let it fail

• Asking for forgiveness is an option using try, except and relying heavily on testing and asserts.
• However, exception handling can get tricky especially if there is no obvious way of resolving something - in this case an explicit function dealing with datasets is more pythonic than a "potential" compatible api that xarray has between dataarray and datasets (explicit > implicit)

Further, most alternatives don't address this in particular:

• reducing code bloat and targetted testing. Only needing to test for datasets is a constraint that allows one to make a lot more assumptions on the compatibility of the code (while also implicitly testing dataarrays).
• currently tests are heavily skewed towards dataarrays.