Skip to content

Commit

Permalink
Refactor and test getting netcdf variable to xr.DataArray
Browse files Browse the repository at this point in the history
  • Loading branch information
lahtinep committed Mar 7, 2024
1 parent f6d87f3 commit f1bc699
Show file tree
Hide file tree
Showing 2 changed files with 75 additions and 11 deletions.
29 changes: 18 additions & 11 deletions satpy/readers/netcdf_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -258,17 +258,7 @@ def collect_cache_vars(self, cache_var_size):
cache_vars = self._collect_cache_var_names(cache_var_size)
for var_name in cache_vars:
v = self.file_content[var_name]
try:
attrs = v.attrs
except AttributeError:
attrs = v.__dict__
try:
arr = xr.DataArray(
v[:], dims=v.dimensions, attrs=attrs, name=v.name)
except (ValueError, IndexError):
# Handle scalars for h5netcdf backend
arr = xr.DataArray(
v.__array__(), dims=v.dimensions, attrs=attrs, name=v.name)
arr = get_data_as_xarray(v)
self.cached_file_content[var_name] = arr

def _collect_cache_var_names(self, cache_var_size):
Expand Down Expand Up @@ -384,6 +374,23 @@ def _compose_replacement_names(variable_name_replacements, var, variable_names):
variable_names.append(var.format(**{key: val}))


def get_data_as_xarray(variable):
"""Get data in variable as xr.DataArray."""
try:
attrs = variable.attrs
except AttributeError:
attrs = variable.__dict__
try:
arr = xr.DataArray(
variable[:], dims=variable.dimensions, attrs=attrs, name=variable.name)
except (ValueError, IndexError):
# Handle scalars for h5netcdf backend
arr = xr.DataArray(
variable.__array__(), dims=variable.dimensions, attrs=attrs, name=variable.name)

return arr


class NetCDF4FsspecFileHandler(NetCDF4FileHandler):
"""NetCDF4 file handler using fsspec to read files remotely."""

Expand Down
57 changes: 57 additions & 0 deletions satpy/tests/reader_tests/test_netcdf_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -293,3 +293,60 @@ def test_use_h5netcdf_for_file_not_accessible_locally(self):
fh = NetCDF4FsspecFileHandler(fname, {}, {})
h5_file.assert_called_once()
assert fh._use_h5netcdf


NC_ATTRS = {
"standard_name": "test_data",
"scale_factor": 0.01,
"add_offset": 0}

def test_get_data_as_xarray_netcdf4():
"""Test getting xr.DataArray from netcdf4 variable."""
import tempfile

import netCDF4 as nc
import numpy as np

from satpy.readers.netcdf_utils import get_data_as_xarray

data = np.array([1, 2, 3])

with tempfile.TemporaryDirectory() as tmpdir:
# Create an empty HDF5
fname = os.path.join(tmpdir, "test.nc")
dset = nc.Dataset(fname, "w")
dset.createDimension("y", None)
var = dset.createVariable("test_data", "uint8", ("y",))
var[:] = data
var.setncatts(NC_ATTRS)
# Turn off automatic scale factor and offset handling
dset.set_auto_maskandscale(False)
res = get_data_as_xarray(var)
np.testing.assert_equal(res.data, data)
assert res.attrs == NC_ATTRS
dset.close()


def test_get_data_as_xarray_h5netcdf():
"""Test getting xr.DataArray from h5netcdf variable."""
import tempfile

import h5netcdf
import numpy as np

from satpy.readers.netcdf_utils import get_data_as_xarray

data = np.array([1, 2, 3])

with tempfile.TemporaryDirectory() as tmpdir:
# Create an empty HDF5
fname = os.path.join(tmpdir, "test.nc")
with h5netcdf.File(fname, "w") as fid:
fid.dimensions = {"y": data.size}
var = fid.create_variable("test_data", ("y",), "uint8")
var[:] = data
for key in NC_ATTRS:
var.attrs[key] = NC_ATTRS[key]
res = get_data_as_xarray(var)
np.testing.assert_equal(res.data, data)
assert res.attrs == NC_ATTRS

0 comments on commit f1bc699

Please sign in to comment.