DataAssimBench/dabench/data/_xarray_accessor.py at 37c50e2befafeb773f774d41d040eec28b8eab64 · StevePny/DataAssimBench · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
import xarray as xr
import numpy as np
import warnings


def _check_split_lengths(
        xr_obj: xr.Dataset | xr.DataArray,
        split_lengths: list | np.ndarray):
    total_length = np.sum(split_lengths)
    xr_timedim = xr_obj.sizes['time']
    if xr_timedim < total_length:
        warnings.warn("Specified split lengths ({}) exceed \n"
                      "Xarray object's time dimension ({}).".format(
                          split_lengths, xr_timedim
                      ))
    elif xr_timedim > total_length:
        warnings.warn("Specified split lengths ({}) are shorter than "
                      "Xarray object's time dimension ({}).".format(
                          split_lengths, xr_timedim
                      ))


@xr.register_dataset_accessor('dab')
class DABenchDatasetAccessor:
    """Helper methods for manipulating xarray Datasets"""
    def __init__(self,
                 xarray_obj: xr.Dataset):
        self._obj = xarray_obj

    def flatten(self) -> xr.DataArray:
        if 'time' in self._obj.coords:
            remaining_dim = ['time']
        else:
            remaining_dim = []
        return self._obj.to_stacked_array('system', remaining_dim)

    def split_train_val_test(self,
                             split_lengths: list | np.ndarray
                             ) -> tuple[xr.Dataset,  ...]:
        if (np.array(split_lengths) > 1.0).sum() == 0:
            # Assuming split_lengths is provided as fraction
            split_lengths = np.round(
                    np.array(split_lengths)*self._obj.sizes['time']
                    ).astype(int)
        _check_split_lengths(self._obj, split_lengths)
        out_ds = []
        start_i = 0
        for sl in split_lengths:
            end_i = start_i + sl
            out_ds.append(self._obj.isel(time=slice(start_i, end_i)))
            start_i = end_i
        return tuple(out_ds)


@xr.register_dataarray_accessor('dab')
class DABenchDataArrayAccessor:
    """Helper methods for manipulating xarray DataArrays"""
    def __init__(self,
                 xarray_obj: xr.DataArray):
        self._obj = xarray_obj

    def unflatten(self) -> xr.Dataset:
        return self._obj.to_unstacked_dataset('system')

    def split_train_val_test(self,
                             split_lengths: list | np.ndarray
                             ) -> tuple[xr.DataArray, ...]:
        if (np.array(split_lengths) > 1.0).sum() == 0:
            # Assuming split_lengths is provided as fraction
            split_lengths = np.round(
                    np.array(split_lengths)*self._obj.sizes['time']
                    ).astype(int)
        _check_split_lengths(self._obj, split_lengths)
        out_ds = []
        start_i = 0
        for sl in split_lengths:
            end_i = start_i + sl
            out_ds.append(self._obj.isel(time=slice(start_i, end_i)))
        return tuple(out_ds)