Coverage for C: \ Users \ peaco \ OneDrive \ Documents \ GitHub \ mth5 \ mth5 \ timeseries \ xarray_helpers.py: 100%
39 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-01-10 00:01 -0800
1"""
2Module containing helper functions for working with xarray objects.
3"""
5from typing import Optional, Union
7import numpy as np
8import xarray as xr
11def covariance_xr(
12 X: xr.DataArray,
13 aweights: Optional[Union[np.ndarray, None]] = None,
14 bias: Optional[bool] = False,
15 rowvar: Optional[bool] = False,
16) -> xr.DataArray:
17 """
18 Compute the covariance matrix with numpy.cov.
20 Parameters
21 ----------
22 X: xarray.core.dataarray.DataArray
23 Multivariate time series as an xarray
24 aweights: array_like, optional
25 Passthrough param for np.cov.
26 1-D array of observation vector weights. These relative weights are
27 typically large for observations considered "important" and smaller for
28 observations considered less "important". If ``ddof=0`` the array of
29 weights can be used to assign probabilities to observation vectors.
30 bias: bool
31 Passthrough param for np.cov.
32 Default normalization (False) is by ``(N - 1)``, where ``N`` is the
33 number of observations given (unbiased estimate). If `bias` is True,
34 then normalization is by ``N``. These values can be overridden by using
35 the keyword ``ddof`` in numpy versions >= 1.5.
36 rowvar: bool
37 Passthrough param for np.cov.
38 If `rowvar` is True (default), then each row represents a
39 variable, with observations in the columns. Otherwise, the relationship
40 is transposed: each column represents a variable, while the rows
41 contain observations.
43 Returns
44 -------
45 S: xarray.DataArray
46 The covariance matrix of the data in xarray form.
48 Development Notes:
49 In case of ValueError: conflicting sizes for dimension 'channel_1', this likely means the bool for rowvar
50 should be flipped.
51 """
53 channels = list(X.coords["variable"].values)
55 cov_result = np.cov(X.values, rowvar=rowvar, aweights=aweights, bias=bias)
57 # Handle scalar result for single variable
58 if cov_result.ndim == 0:
59 cov_result = np.array([[cov_result]])
61 S = xr.DataArray(
62 cov_result,
63 dims=["channel_1", "channel_2"],
64 coords={"channel_1": channels, "channel_2": channels},
65 )
66 return S
69def initialize_xrda_1d(
70 channels: list,
71 dtype: Optional[type] = None,
72 value: Optional[Union[complex, float, bool]] = 0.0,
73) -> xr.DataArray:
74 """
75 Returns a 1D xr.DataArray with variable "channel", having values channels named by the input list.
77 Parameters
78 ----------
79 channels: list
80 The channels in the multivariate array
81 dtype: type, optional
82 The datatype to initialize the array.
83 Common cases are complex, float, and bool
84 value: Union[complex, float, bool], optional
85 The default value to assign the array
87 Returns
88 -------
89 xrda: xarray.core.dataarray.DataArray
90 An xarray container for the channels, initialized to zeros.
91 """
92 k = len(channels)
93 xrda = xr.DataArray(
94 np.zeros(k, dtype=dtype),
95 dims=[
96 "variable",
97 ],
98 coords={
99 "variable": channels,
100 },
101 )
102 if value != 0:
103 data = value * np.ones(k, dtype=dtype)
104 xrda.data = data
105 return xrda
108def initialize_xrds_2d(
109 variables: list,
110 coords: dict,
111 dtype: Optional[type] = complex,
112 value: Optional[Union[complex, float, bool]] = 0,
113) -> xr.Dataset:
114 """
115 Returns a 2D xr.Dataset with the given variables and coordinates.
117 Parameters
118 ----------
119 variables: list
120 List of variable names to create in the dataset
121 coords: dict
122 Dictionary of coordinates for the dataset dimensions
123 dtype: type, optional
124 The datatype to initialize the arrays.
125 Common cases are complex, float, and bool
126 value: Union[complex, float, bool], optional
127 The default value to assign the arrays
129 Returns
130 -------
131 xrds: xr.Dataset
132 A 2D xarray Dataset with dimensions from coords
133 """
134 # Get dimensions from coords
135 dims = list(coords.keys())
136 shape = tuple(len(v) for v in coords.values())
138 # Initialize empty dataset
139 xrds = xr.Dataset(coords=coords)
141 # Add each variable
142 for var in variables:
143 if value == 0:
144 data = np.zeros(shape, dtype=dtype)
145 else:
146 data = value * np.ones(shape, dtype=dtype)
148 xrds[var] = xr.DataArray(data, dims=dims, coords=coords)
150 return xrds
153def initialize_xrda_2d(
154 variables: list, coords: dict, dtype: type = complex, value: Union[int, float] = 0.0
155) -> xr.Dataset:
156 """Initialize a 3D xarray DataArray with dimensions from coords plus 'variable'.
158 Parameters
159 ----------
160 variables : list
161 List of variable names for the additional dimension.
162 coords : dict
163 Dictionary of coordinates for the dataset dimensions.
164 dtype : type, optional
165 Data type for the array, by default complex.
166 value : int or float, optional
167 Value to initialize the array with, by default 0.
169 Returns
170 -------
171 xr.DataArray
172 A 3D DataArray with dimensions from coords plus 'variable'.
173 """
174 # Create Dataset first
175 ds = initialize_xrds_2d(variables, coords, dtype, value)
177 # Convert to DataArray with original dimension order plus 'variable'
178 dims = list(coords.keys())
179 da = ds.to_array(dim="variable").transpose(*dims, "variable")
181 return da
184def initialize_xrda_2d_cov(
185 channels, dtype=complex, value: Optional[Union[complex, float, bool]] = 0
186):
187 """
188 TODO: consider changing nomenclature from dims=["channel_1", "channel_2"],
189 to dims=["variable_1", "variable_2"], to be consistent with initialize_xrda_1d
191 Parameters
192 ----------
193 channels: list
194 The channels in the multivariate array. The covariance matrix will be square
195 with dimensions len(channels) x len(channels).
196 dtype: type
197 The datatype to initialize the array.
198 Common cases are complex, float, and bool
199 value: Union[complex, float, bool]
200 The default value to assign the array
202 Returns
203 -------
204 xrda: xarray.core.dataarray.DataArray
205 An xarray container for the channel covariances, initialized to zeros.
206 The matrix is square with dimensions len(channels) x len(channels).
207 """
208 K = len(channels)
209 xrda = xr.DataArray(
210 np.zeros((K, K), dtype=dtype),
211 dims=["channel_1", "channel_2"],
212 coords={"channel_1": channels, "channel_2": channels},
213 )
214 if value != 0:
215 data = value * np.ones(xrda.shape, dtype=dtype)
216 xrda.data = data
218 return xrda