-
Notifications
You must be signed in to change notification settings - Fork 15
/
xarray.py
66 lines (47 loc) · 1.65 KB
/
xarray.py
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
"""Xarray integration.
- conversion between 1D histograms and xarray Datasets.
This is experimental and may change.
"""
from __future__ import annotations
from typing import TYPE_CHECKING
from xarray import DataArray, Dataset
from physt.types import Histogram1D
if TYPE_CHECKING:
from typing import Any, Dict, Type
def _h1_to_xarray(h1: Histogram1D) -> Dataset:
"""Convert histogram to an xarray Datates representation.
See also: Histogram1D.from_xarray (inverse operation)
"""
# TODO: Rethink how the dimensions / variables are used
data_vars: Dict[str, Any] = {
"frequencies": DataArray(h1.frequencies, dims="bin"),
"errors2": DataArray(h1.errors2, dims="bin"),
"bins": DataArray(h1.bins, dims=("bin", "x01")),
}
coords: Dict[str, Any] = {}
attrs: Dict[str, Any] = {
"underflow": h1.underflow,
"overflow": h1.overflow,
"inner_missed": h1.inner_missed,
"keep_missed": h1.keep_missed,
}
attrs.update(h1.meta_data)
# TODO: Add stats
return Dataset(data_vars, coords, attrs)
def _h1_from_xarray(cls: Type[Histogram1D], arr: Dataset) -> Histogram1D:
"""Convert form xarray.Dataset
Parameters
----------
arr: The data in xarray representation
"""
kwargs = {
"frequencies": arr["frequencies"],
"binning": arr["bins"],
"errors2": arr["errors2"],
}
kwargs.update(arr.attrs) # type: ignore
# TODO: Add stats
return cls(**kwargs) # type: ignore
setattr(Histogram1D, "to_xarray", _h1_to_xarray)
setattr(Histogram1D, "from_xarray", classmethod(_h1_from_xarray))
# TODO: Implement multi-dimensional histograms