From 2745f7b8d0589a6f7a3ac4c703fc4226be909191 Mon Sep 17 00:00:00 2001 From: Alexandra Udaltsova <43303448+AUdaltsova@users.noreply.github.com> Date: Fri, 30 Aug 2024 16:39:52 +0100 Subject: [PATCH] Add docstring and change func name --- ocf_data_sampler/load/utils.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/ocf_data_sampler/load/utils.py b/ocf_data_sampler/load/utils.py index e837cd6..7d16610 100644 --- a/ocf_data_sampler/load/utils.py +++ b/ocf_data_sampler/load/utils.py @@ -29,7 +29,15 @@ def make_spatial_coords_increasing(ds: xr.Dataset, x_coord: str, y_coord: str) - return ds -def underlying_array(ds: xr.Dataset, datavar: str | None = None) -> xr.DataArray: +def get_xr_data_array_from_xr_dataset(ds: xr.Dataset, datavar: str | None = None) -> xr.DataArray: + """Return underlying xr.DataArray from passed xr.Dataset. If davar provided, returns specified variable. + If datavar not provided, checks only one variable is present and returns it as an xr.DataArray. + + Args: + ds: xr.Dataset to extract xr.DataArray from + datavar: (if provided) return specific variable from ds + """ + if datavar: return ds[datavar] datavars = list(ds.var())