-"""
-
-Credits:
-
-This file was adopted from: https://github.com/pydata/xarray # noqa
-Source file: https://github.com/pydata/xarray/blob/1d7bcbdc75b6d556c04e2c7d7a042e4379e15303/xarray/backends/rasterio_.py # noqa
-"""
-# pylint: disable=too-many-lines
-import contextlib
-import functools
-import importlib.metadata
-import os
-import re
-import threading
-import warnings
-from collections import defaultdict
-from collections.abc import Hashable, Iterable
-from typing import Any, Optional, Union
-
-import numpy
-import rasterio
-from numpy.typing import NDArray
-from packaging import version
-from rasterio.errors import NotGeoreferencedWarning
-from rasterio.vrt import WarpedVRT
-from xarray import Dataset, IndexVariable
-from xarray.backends.common import BackendArray
-from xarray.backends.file_manager import CachingFileManager, FileManager
-from xarray.backends.locks import SerializableLock
-from xarray.coding import times, variables
-from xarray.core import indexing
-from xarray.core.dataarray import DataArray
-from xarray.core.dtypes import maybe_promote
-from xarray.core.utils import is_scalar
-from xarray.core.variable import as_variable
-
-from rioxarray.exceptions import RioXarrayError
-from rioxarray.rioxarray import _generate_spatial_coords
-
-FILL_VALUE_NAMES = ("_FillValue", "missing_value", "fill_value", "nodata")
-UNWANTED_RIO_ATTRS = ("nodatavals", "is_tiled", "res")
-# TODO: should this be GDAL_LOCK instead?
-RASTERIO_LOCK = SerializableLock()
-NO_LOCK = contextlib.nullcontext()
-
-
-def _ensure_warped_vrt(riods, vrt_params):
- """
- Ensuire the dataset is represented as a warped vrt
- """
- if vrt_params is None:
- return riods
- if isinstance(riods, SingleBandDatasetReader):
- riods._create_vrt(vrt_params)
- else:
- riods = WarpedVRT(riods, **vrt_params)
- return riods
-
-
-class SingleBandDatasetReader:
- """
- Hack to have a DatasetReader behave like it only has one band
- """
-
- def __init__(self, riods, bidx, vrt_params=None) -> None:
- self._riods = riods
- self._bidx = bidx
- self._vrt_params = vrt_params
- self._create_vrt(vrt_params=vrt_params)
-
- def __getattr__(self, __name: str) -> Any:
- return getattr(self._riods, __name)
-
- def _create_vrt(self, vrt_params):
- if vrt_params is not None and not isinstance(self._riods, WarpedVRT):
- self._riods = WarpedVRT(self._riods, **vrt_params)
- self._vrt_params = vrt_params
-
- @property
- def name(self):
- """
- str: name of the dataset. Usually the path.
- """
- if isinstance(self._riods, rasterio.vrt.WarpedVRT):
- return self._riods.src_dataset.name
- return self._riods.name
-
- @property
- def count(self):
- """
- int: band count
- """
- return 1
-
- @property
- def nodata(self):
- """
- Nodata value for the band
- """
- return self._riods.nodatavals[self._bidx]
-
- @property
- def offsets(self):
- """
- Offset value for the band
- """
- return [self._riods.offsets[self._bidx]]
-
- @property
- def scales(self):
- """
- Scale value for the band
- """
- return [self._riods.scales[self._bidx]]
-
- @property
- def units(self):
- """
- Unit for the band
- """
- return [self._riods.units[self._bidx]]
-
- @property
- def descriptions(self):
- """
- Description for the band
- """
- return [self._riods.descriptions[self._bidx]]
-
- @property
- def dtypes(self):
- """
- dtype for the band
- """
- return [self._riods.dtypes[self._bidx]]
-
- @property
- def indexes(self):
- """
- indexes for the band
- """
- return [self._riods.indexes[self._bidx]]
-
- def read(self, indexes=None, **kwargs): # pylint: disable=unused-argument
- """
- read data for the band
- """
- return self._riods.read(indexes=self._bidx + 1, **kwargs)
-
- def tags(self, bidx=None, **kwargs): # pylint: disable=unused-argument
- """
- read tags for the band
- """
- return self._riods.tags(bidx=self._bidx + 1, **kwargs)
-
-
-RasterioReader = Union[
- rasterio.io.DatasetReader, rasterio.vrt.WarpedVRT, SingleBandDatasetReader
-]
-
-
-try:
- _DASK_GTE_018 = version.parse(importlib.metadata.version("dask")) >= version.parse(
- "0.18.0"
- )
-except importlib.metadata.PackageNotFoundError:
- _DASK_GTE_018 = False
-
-
-def _get_unsigned_dtype(unsigned, dtype):
- """
- Based on: https://github.com/pydata/xarray/blob/abe1e613a96b000ae603c53d135828df532b952e/xarray/coding/variables.py#L306-L334
- """
- dtype = numpy.dtype(dtype)
- if unsigned is True and dtype.kind == "i":
- return numpy.dtype(f"u{dtype.itemsize}")
- if unsigned is False and dtype.kind == "u":
- return numpy.dtype(f"i{dtype.itemsize}")
- return None
-
-
-class FileHandleLocal(threading.local):
- """
- This contains the thread local ThreadURIManager
- """
-
- def __init__(self): # pylint: disable=super-init-not-called
- self.thread_manager = None # Initialises in each thread
-
-
-class ThreadURIManager:
- """
- This handles opening & closing file handles in each thread.
- """
-
- def __init__(
- self,
- opener,
- *args,
- mode="r",
- kwargs=None,
- ):
- self._opener = opener
- self._args = args
- self._mode = mode
- self._kwargs = {} if kwargs is None else dict(kwargs)
- self._file_handle = None
-
- @property
- def file_handle(self):
- """
- File handle returned by the opener.
- """
- if self._file_handle is not None:
- return self._file_handle
- self._file_handle = self._opener(*self._args, mode=self._mode, **self._kwargs)
- return self._file_handle
-
- def close(self):
- """
- Close file handle.
- """
- if self._file_handle is not None:
- self._file_handle.close()
- self._file_handle = None
-
- def __del__(self):
- self.close()
-
- def __enter__(self):
- return self
-
- def __exit__(self, type_, value, traceback):
- self.close()
-
-
-class URIManager(FileManager):
- """
- The URI manager is used for lockless reading
- """
-
- def __init__(
- self,
- opener,
- *args,
- mode="r",
- kwargs=None,
- ):
- self._opener = opener
- self._args = args
- self._mode = mode
- self._kwargs = {} if kwargs is None else dict(kwargs)
- self._local = FileHandleLocal()
-
- def acquire(self, needs_lock=True):
- if self._local.thread_manager is None:
- self._local.thread_manager = ThreadURIManager(
- self._opener, *self._args, mode=self._mode, kwargs=self._kwargs
- )
- return self._local.thread_manager.file_handle
-
- @contextlib.contextmanager
- def acquire_context(self, needs_lock=True):
- try:
- yield self.acquire(needs_lock=needs_lock)
- except Exception:
- self.close(needs_lock=needs_lock)
- raise
-
- def close(self, needs_lock=True):
- if self._local.thread_manager is not None:
- self._local.thread_manager.close()
- self._local.thread_manager = None
-
- def __del__(self):
- self.close(needs_lock=False)
-
- def __getstate__(self):
- """State for pickling."""
- return (self._opener, self._args, self._mode, self._kwargs)
-
- def __setstate__(self, state):
- """Restore from a pickle."""
- opener, args, mode, kwargs = state
- self.__init__(opener, *args, mode=mode, kwargs=kwargs)
-
-
-class RasterioArrayWrapper(BackendArray):
- """A wrapper around rasterio dataset objects"""
-
- # pylint: disable=too-many-instance-attributes
-
- def __init__(
- self,
- manager,
- lock,
- name,
- vrt_params=None,
- masked=False,
- mask_and_scale=False,
- unsigned=False,
- ):
- self.manager = manager
- self.lock = lock
- self.masked = masked or mask_and_scale
- self.mask_and_scale = mask_and_scale
-
- # cannot save riods as an attribute: this would break pickleability
- riods = _ensure_warped_vrt(manager.acquire(), vrt_params)
- self.vrt_params = vrt_params
- self._shape = (riods.count, riods.height, riods.width)
-
- self._dtype = None
- self._unsigned_dtype = None
- self._fill_value = riods.nodata
- dtypes = riods.dtypes
- if not numpy.all(numpy.asarray(dtypes) == dtypes[0]):
- raise ValueError("All bands should have the same dtype")
-
- dtype = _rasterio_to_numpy_dtype(dtypes)
- if mask_and_scale and unsigned is not None:
- self._unsigned_dtype = _get_unsigned_dtype(
- unsigned=unsigned,
- dtype=dtype,
- )
- if self._unsigned_dtype is not None and self._fill_value is not None:
- self._fill_value = self._unsigned_dtype.type(self._fill_value)
- if self._unsigned_dtype is None and dtype.kind not in ("i", "u"):
- warnings.warn(
- f"variable {name!r} has _Unsigned attribute but is not "
- "of integer type. Ignoring attribute.",
- variables.SerializationWarning,
- stacklevel=3,
- )
- if self.masked:
- self._dtype, self._fill_value = maybe_promote(dtype)
- else:
- self._dtype = dtype
-
- @property
- def dtype(self):
- """
- Data type of the array
- """
- return self._dtype
-
- @property
- def fill_value(self):
- """
- Fill value of the array
- """
- return self._fill_value
-
- @property
- def shape(self):
- """
- Shape of the array
- """
- return self._shape
-
- def _get_indexer(self, key):
- """Get indexer for rasterio array.
-
- Parameter
- ---------
- key: tuple of int
-
- Returns
- -------
- band_key: an indexer for the 1st dimension
- window: two tuples. Each consists of (start, stop).
- squeeze_axis: axes to be squeezed
- np_ind: indexer for loaded numpy array
-
- See also
- --------
- indexing.decompose_indexer
- """
- if len(key) != 3:
- raise RioXarrayError("rasterio datasets should always be 3D")
-
- # bands cannot be windowed but they can be listed
- band_key = key[0]
- np_inds = []
- # bands (axis=0) cannot be windowed but they can be listed
- if isinstance(band_key, slice):
- start, stop, step = band_key.indices(self.shape[0])
- band_key = numpy.arange(start, stop, step)
- # be sure we give out a list
- band_key = (numpy.asarray(band_key) + 1).tolist()
- if isinstance(band_key, list): # if band_key is not a scalar
- np_inds.append(slice(None))
-
- # but other dims can only be windowed
- window = []
- squeeze_axis = []
- for iii, (ikey, size) in enumerate(zip(key[1:], self.shape[1:])):
- if isinstance(ikey, slice):
- # step is always positive. see indexing.decompose_indexer
- start, stop, step = ikey.indices(size)
- np_inds.append(slice(None, None, step))
- elif is_scalar(ikey):
- # windowed operations will always return an array
- # we will have to squeeze it later
- squeeze_axis.append(-(2 - iii))
- start = ikey
- stop = ikey + 1
- else:
- start, stop = numpy.min(ikey), numpy.max(ikey) + 1
- np_inds.append(ikey - start)
- window.append((start, stop))
-
- if isinstance(key[1], numpy.ndarray) and isinstance(key[2], numpy.ndarray):
- # do outer-style indexing
- np_inds[-2:] = numpy.ix_(*np_inds[-2:])
-
- return band_key, tuple(window), tuple(squeeze_axis), tuple(np_inds)
-
- def _getitem(self, key):
- band_key, window, squeeze_axis, np_inds = self._get_indexer(key)
- if not band_key or any(start == stop for (start, stop) in window):
- # no need to do IO
- shape = (len(band_key),) + tuple(stop - start for (start, stop) in window)
- out = numpy.zeros(shape, dtype=self.dtype)
- else:
- with self.lock:
- riods = _ensure_warped_vrt(
- self.manager.acquire(needs_lock=False), self.vrt_params
- )
- out = riods.read(band_key, window=window, masked=self.masked)
- if self._unsigned_dtype is not None:
- out = out.astype(self._unsigned_dtype)
- if self.masked:
- out = numpy.ma.filled(out.astype(self.dtype), self.fill_value)
- if self.mask_and_scale:
- if not isinstance(band_key, Iterable):
- out = (
- out * riods.scales[band_key - 1]
- + riods.offsets[band_key - 1]
- )
- else:
- for iii, band_iii in enumerate(numpy.atleast_1d(band_key) - 1):
- out[iii] = (
- out[iii] * riods.scales[band_iii]
- + riods.offsets[band_iii]
- )
-
- if squeeze_axis:
- out = numpy.squeeze(out, axis=squeeze_axis)
- return out[np_inds]
-
- def __getitem__(self, key):
- return indexing.explicit_indexing_adapter(
- key, self.shape, indexing.IndexingSupport.OUTER, self._getitem
- )
-
-
-def _parse_envi(meta):
- """Parse ENVI metadata into Python data structures.
-
- See the link for information on the ENVI header file format:
- http://www.harrisgeospatial.com/docs/enviheaderfiles.html
-
- Parameters
- ----------
- meta : dict
- Dictionary of keys and str values to parse, as returned by the rasterio
- tags(ns='ENVI') call.
-
- Returns
- -------
- parsed_meta : dict
- Dictionary containing the original keys and the parsed values
-
- """
-
- def parsevec(value):
- return numpy.fromstring(value.strip("{}"), dtype="float", sep=",")
-
- def default(value):
- return value.strip("{}")
-
- parse = {"wavelength": parsevec, "fwhm": parsevec}
- parsed_meta = {key: parse.get(key, default)(value) for key, value in meta.items()}
- return parsed_meta
-
-
-def _rasterio_to_numpy_dtype(dtypes):
- """Numpy dtype from first entry of rasterio dataset.dtypes"""
- # rasterio has some special dtype names (complex_int16 -> numpy.complex64)
- if dtypes[0] == "complex_int16":
- dtype = numpy.dtype("complex64")
- else:
- dtype = numpy.dtype(dtypes[0])
-
- return dtype
-
-
-def _to_numeric(value: Any) -> float:
- """
- Convert the value to a number
- """
- try:
- value = int(value)
- except (TypeError, ValueError):
- try:
- value = float(value)
- except (TypeError, ValueError):
- pass
- return value
-
-
-def _parse_tag(key: str, value: Any) -> tuple[str, Any]:
- # NC_GLOBAL is appended to tags with netcdf driver and is not really needed
- key = key.split("NC_GLOBAL#")[-1]
- if value.startswith("{") and value.endswith("}"):
- try:
- new_val = numpy.fromstring(value.strip("{}"), dtype="float", sep=",")
- # pylint: disable=len-as-condition
- value = new_val if len(new_val) else _to_numeric(value)
- except ValueError:
- value = _to_numeric(value)
- else:
- value = _to_numeric(value)
- return key, value
-
-
-def _parse_tags(tags: dict) -> dict:
- parsed_tags = {}
- for key, value in tags.items():
- key, value = _parse_tag(key, value)
- parsed_tags[key] = value
- return parsed_tags
-
-
-NETCDF_DTYPE_MAP = {
- 0: object, # NC_NAT
- 1: numpy.byte, # NC_BYTE
- 2: numpy.char, # NC_CHAR
- 3: numpy.short, # NC_SHORT
- 4: numpy.int_, # NC_INT, NC_LONG
- 5: float, # NC_FLOAT
- 6: numpy.double, # NC_DOUBLE
- 7: numpy.ubyte, # NC_UBYTE
- 8: numpy.ushort, # NC_USHORT
- 9: numpy.uint, # NC_UINT
- 10: numpy.int64, # NC_INT64
- 11: numpy.uint64, # NC_UINT64
- 12: object, # NC_STRING
-}
-
-
-def _load_netcdf_attrs(tags: dict, data_array: DataArray) -> None:
- """
- Loads the netCDF attributes into the data array
-
- Attributes stored in this format:
- - variable_name#attr_name: attr_value
- """
- for key, value in tags.items():
- key, value = _parse_tag(key, value)
- key_split = key.split("#")
- if len(key_split) != 2:
- continue
- variable_name, attr_name = key_split
- if variable_name in data_array.coords:
- data_array.coords[variable_name].attrs.update({attr_name: value})
-
-
-def _parse_netcdf_attr_array(attr: Union[NDArray, str], dtype=None) -> NDArray:
- """
- Expected format: '{2,6}' or '[2. 6.]'
- """
- value: Union[NDArray, str, list]
- if isinstance(attr, str):
- if attr.startswith("{"):
- value = attr.strip("{}").split(",")
- else:
- value = attr.strip("[]").split()
- elif not isinstance(attr, Iterable):
- value = [attr]
- else:
- value = attr
- return numpy.array(value, dtype=dtype)
-
-
-def _load_netcdf_1d_coords(tags: dict) -> dict:
- """
- Dimension information:
- - NETCDF_DIM_EXTRA: '{time}' (comma separated list of dim names)
- - NETCDF_DIM_time_DEF: '{2,6}' or '[2. 6.]' (dim size, dim dtype)
- - NETCDF_DIM_time_VALUES: '{0,872712.659688}' (comma separated list of data) or [ 0. 872712.659688]
- """
- dim_names = tags.get("NETCDF_DIM_EXTRA")
- if not dim_names:
- return {}
- dim_names = _parse_netcdf_attr_array(dim_names)
- coords = {}
- for dim_name in dim_names:
- dim_def = tags.get(f"NETCDF_DIM_{dim_name}_DEF")
- if dim_def is None:
- continue
- # pylint: disable=unused-variable
- dim_size, dim_dtype = _parse_netcdf_attr_array(dim_def)
- dim_dtype = NETCDF_DTYPE_MAP.get(int(float(dim_dtype)), object)
- dim_values = _parse_netcdf_attr_array(tags[f"NETCDF_DIM_{dim_name}_VALUES"])
- coords[dim_name] = IndexVariable(dim_name, dim_values)
- return coords
-
-
-def build_subdataset_filter(
- group_names: Optional[Union[str, list[str], tuple[str, ...]]],
- variable_names: Optional[Union[str, list[str], tuple[str, ...]]],
-):
- """
- Example::
- 'HDF4_EOS:EOS_GRID:"./modis/MOD09GQ.A2017290.h11v04.006.NRT.hdf":
- MODIS_Grid_2D:sur_refl_b01_1'
-
- Parameters
- ----------
- group_names: str or list or tuple
- Name or names of netCDF groups to filter by.
-
- variable_names: str or list or tuple
- Name or names of netCDF variables to filter by.
-
- Returns
- -------
- re.SRE_Pattern: output of re.compile()
- """
- variable_query = r"\w+"
- if variable_names is not None:
- if not isinstance(variable_names, (tuple, list)):
- variable_names = [variable_names]
- variable_names = [re.escape(variable_name) for variable_name in variable_names]
- variable_query = rf"(?:{'|'.join(variable_names)})"
- if group_names is not None:
- if not isinstance(group_names, (tuple, list)):
- group_names = [group_names]
- group_names = [re.escape(group_name) for group_name in group_names]
- group_query = rf"(?:{'|'.join(group_names)})"
- else:
- return re.compile(r"".join([r".*(?:\:/|\:)(/+)?", variable_query, r"$"]))
- return re.compile(
- r"".join(
- [r".*(?:\:/|\:)(/+)?", group_query, r"[:/](/+)?", variable_query, r"$"]
- )
- )
-
-
-def _get_rasterio_attrs(riods: RasterioReader):
- """
- Get rasterio specific attributes
- """
- # pylint: disable=too-many-branches
- # Add rasterio attributes
- attrs = _parse_tags({**riods.tags(), **riods.tags(1)})
- # remove attributes with informaiton
- # that should be added by GDAL/rasterio
- for unwanted_attr in FILL_VALUE_NAMES + UNWANTED_RIO_ATTRS:
- attrs.pop(unwanted_attr, None)
- if riods.nodata is not None:
- # The nodata values for the raster bands
- attrs["_FillValue"] = riods.nodata
- # The scale values for the raster bands
- if len(set(riods.scales)) > 1:
- attrs["scales"] = riods.scales
- warnings.warn(
- "Offsets differ across bands. The 'scale_factor' attribute will "
- "not be added. See the 'scales' attribute."
- )
- else:
- attrs["scale_factor"] = riods.scales[0]
- # The offset values for the raster bands
- if len(set(riods.offsets)) > 1:
- attrs["offsets"] = riods.offsets
- warnings.warn(
- "Offsets differ across bands. The 'add_offset' attribute will "
- "not be added. See the 'offsets' attribute."
- )
- else:
- attrs["add_offset"] = riods.offsets[0]
- if any(riods.descriptions):
- if len(set(riods.descriptions)) == 1:
- attrs["long_name"] = riods.descriptions[0]
- else:
- # Descriptions for each dataset band
- attrs["long_name"] = riods.descriptions
- if any(riods.units):
- # A list of units string for each dataset band
- if len(riods.units) == 1:
- attrs["units"] = riods.units[0]
- else:
- attrs["units"] = riods.units
-
- return attrs
-
-
-def _decode_datetime_cf(
- data_array: DataArray,
- decode_times: bool,
- decode_timedelta: Optional[bool],
-) -> DataArray:
- """
- Decide the datetime based on CF conventions
- """
- if decode_timedelta is None:
- decode_timedelta = decode_times
-
- for coord in data_array.coords:
- time_var = None
- if decode_times and "since" in data_array[coord].attrs.get("units", ""):
- time_var = times.CFDatetimeCoder(use_cftime=True).decode(
- as_variable(data_array[coord]), name=coord
- )
- elif (
- decode_timedelta
- and data_array[coord].attrs.get("units") in times.TIME_UNITS
- ):
- time_var = times.CFTimedeltaCoder().decode(
- as_variable(data_array[coord]), name=coord
- )
- if time_var is not None:
- dimensions, data, attributes, encoding = variables.unpack_for_decoding(
- time_var
- )
- data_array = data_array.assign_coords(
- {
- coord: IndexVariable(
- dims=dimensions,
- data=data,
- attrs=attributes,
- encoding=encoding,
- )
- }
- )
- return data_array
-
-
-def _parse_driver_tags(
- riods: RasterioReader,
- attrs: dict,
- coords: dict,
-) -> None:
- # Parse extra metadata from tags, if supported
- parsers = {"ENVI": _parse_envi}
-
- driver = riods.driver
- if driver in parsers:
- meta = parsers[driver](riods.tags(ns=driver))
-
- for key, value in meta.items():
- # Add values as coordinates if they match the band count,
- # as attributes otherwise
- if isinstance(value, (list, numpy.ndarray)) and len(value) == riods.count:
- coords[key] = ("band", numpy.asarray(value))
- else:
- attrs[key] = value
-
-
-def _pop_global_netcdf_attrs_from_vars(dataset_to_clean: Dataset) -> Dataset:
- # remove GLOBAL netCDF attributes from dataset variables
- for coord in dataset_to_clean.coords:
- for variable in dataset_to_clean.variables:
- dataset_to_clean[variable].attrs = {
- attr: value
- for attr, value in dataset_to_clean[variable].attrs.items()
- if attr not in dataset_to_clean.attrs
- and not attr.startswith(f"{coord}#")
- }
- return dataset_to_clean
-
-
-def _subdataset_groups_to_dataset(
- dim_groups: dict[Hashable, dict[Hashable, DataArray]], global_tags: dict
-) -> Union[Dataset, list[Dataset]]:
- if dim_groups:
- dataset: Union[Dataset, list[Dataset]] = []
- for dim_group in dim_groups.values():
- dataset_group = _pop_global_netcdf_attrs_from_vars(
- Dataset(dim_group, attrs=global_tags)
- )
-
- def _ds_close():
- # pylint: disable=cell-var-from-loop
- for data_var in dim_group.values():
- data_var.close()
-
- dataset_group.set_close(_ds_close)
- dataset.append(dataset_group)
- if len(dataset) == 1:
- dataset = dataset.pop()
- else:
- dataset = Dataset(attrs=global_tags)
- return dataset
-
-
-def _load_subdatasets(
- riods: RasterioReader,
- group: Optional[Union[str, list[str], tuple[str, ...]]],
- variable: Optional[Union[str, list[str], tuple[str, ...]]],
- parse_coordinates: bool,
- chunks: Optional[Union[int, tuple, dict]],
- cache: Optional[bool],
- lock: Any,
- masked: bool,
- mask_and_scale: bool,
- decode_times: bool,
- decode_timedelta: Optional[bool],
- **open_kwargs,
-) -> Union[Dataset, list[Dataset]]:
- """
- Load in rasterio subdatasets
- """
- dim_groups: dict[Hashable, dict[Hashable, DataArray]] = defaultdict(dict)
- subdataset_filter = None
- if any((group, variable)):
- subdataset_filter = build_subdataset_filter(group, variable)
- for subdataset in riods.subdatasets:
- if subdataset_filter is not None and not subdataset_filter.match(subdataset):
- continue
- with rasterio.open(subdataset) as rds:
- shape = rds.shape
- rioda: DataArray = open_rasterio( # type: ignore
- subdataset,
- parse_coordinates=shape not in dim_groups and parse_coordinates,
- chunks=chunks,
- cache=cache,
- lock=lock,
- masked=masked,
- mask_and_scale=mask_and_scale,
- default_name=subdataset.split(":")[-1].lstrip("/").replace("/", "_"),
- decode_times=decode_times,
- decode_timedelta=decode_timedelta,
- **open_kwargs,
- )
- dim_groups[shape][rioda.name] = rioda
- return _subdataset_groups_to_dataset(
- dim_groups=dim_groups, global_tags=_parse_tags(riods.tags())
- )
-
-
-def _load_bands_as_variables(
- riods: RasterioReader,
- parse_coordinates: bool,
- chunks: Optional[Union[int, tuple, dict]],
- cache: Optional[bool],
- lock: Any,
- masked: bool,
- mask_and_scale: bool,
- decode_times: bool,
- decode_timedelta: Optional[bool],
- vrt_params: Optional[dict],
- **open_kwargs,
-) -> Union[Dataset, list[Dataset]]:
- """
- Load in rasterio bands as variables
- """
- global_tags = _parse_tags(riods.tags())
- data_vars = {}
- for band in riods.indexes:
- band_riods = SingleBandDatasetReader(
- riods=riods,
- bidx=band - 1,
- vrt_params=vrt_params,
- )
- band_name = f"band_{band}"
- data_vars[band_name] = (
- open_rasterio( # type: ignore
- band_riods,
- parse_coordinates=band == 1 and parse_coordinates,
- chunks=chunks,
- cache=cache,
- lock=lock,
- masked=masked,
- mask_and_scale=mask_and_scale,
- default_name=band_name,
- decode_times=decode_times,
- decode_timedelta=decode_timedelta,
- **open_kwargs,
- )
- .squeeze() # type: ignore
- .drop("band") # type: ignore
- )
- dataset = Dataset(data_vars, attrs=global_tags)
-
- def _ds_close():
- for data_var in data_vars.values():
- data_var.close()
-
- dataset.set_close(_ds_close)
- return dataset
-
-
-def _prepare_dask(
- result: DataArray,
- riods: RasterioReader,
- filename: Union[str, os.PathLike],
- chunks: Union[int, tuple, dict],
-) -> DataArray:
- """
- Prepare the data for dask computations
- """
- # pylint: disable=import-outside-toplevel
- from dask.base import tokenize
-
- # augment the token with the file modification time
- try:
- mtime = os.path.getmtime(filename)
- except (TypeError, OSError):
- # the filename is probably an s3 bucket rather than a regular file
- mtime = None
-
- if chunks in (True, "auto"):
- from dask.array.core import normalize_chunks
-
- if not _DASK_GTE_018:
- raise NotImplementedError("Automatic chunking requires dask >= 0.18.0")
- block_shape = (1,) + riods.block_shapes[0]
- chunks = normalize_chunks(
- chunks=(1, "auto", "auto"),
- shape=(riods.count, riods.height, riods.width),
- dtype=_rasterio_to_numpy_dtype(riods.dtypes),
- previous_chunks=tuple((c,) for c in block_shape),
- )
- token = tokenize(filename, mtime, chunks)
- name_prefix = f"open_rasterio-{token}"
- return result.chunk(chunks, name_prefix=name_prefix, token=token)
-
-
-def _handle_encoding(
- result: DataArray,
- mask_and_scale: bool,
- masked: bool,
- da_name: Optional[Hashable],
- unsigned: Union[bool, None],
-) -> None:
- """
- Make sure encoding handled properly
- """
- if "grid_mapping" in result.attrs:
- variables.pop_to(result.attrs, result.encoding, "grid_mapping", name=da_name)
- if mask_and_scale:
- if "scale_factor" in result.attrs:
- variables.pop_to(
- result.attrs, result.encoding, "scale_factor", name=da_name
- )
- if "add_offset" in result.attrs:
- variables.pop_to(result.attrs, result.encoding, "add_offset", name=da_name)
- if masked:
- if "_FillValue" in result.attrs:
- variables.pop_to(result.attrs, result.encoding, "_FillValue", name=da_name)
- if "missing_value" in result.attrs:
- variables.pop_to(
- result.attrs, result.encoding, "missing_value", name=da_name
- )
-
- if mask_and_scale and unsigned is not None and "_FillValue" in result.encoding:
- unsigned_dtype = _get_unsigned_dtype(
- unsigned=unsigned,
- dtype=result.encoding["dtype"],
- )
- if unsigned_dtype is not None:
- result.encoding["_FillValue"] = unsigned_dtype.type(
- result.encoding["_FillValue"]
- )
-
-
-def _single_band_open(*args, bidx=0, **kwargs):
- """
- Open file as if it only has a single band
- """
- return SingleBandDatasetReader(
- riods=rasterio.open(*args, **kwargs),
- bidx=bidx,
- )
-
-
-[docs]def open_rasterio(
-
filename: Union[
-
str,
-
os.PathLike,
-
rasterio.io.DatasetReader,
-
rasterio.vrt.WarpedVRT,
-
SingleBandDatasetReader,
-
],
-
parse_coordinates: Optional[bool] = None,
-
chunks: Optional[Union[int, tuple, dict]] = None,
-
cache: Optional[bool] = None,
-
lock: Optional[Any] = None,
-
masked: bool = False,
-
mask_and_scale: bool = False,
-
variable: Optional[Union[str, list[str], tuple[str, ...]]] = None,
-
group: Optional[Union[str, list[str], tuple[str, ...]]] = None,
-
default_name: Optional[str] = None,
-
decode_times: bool = True,
-
decode_timedelta: Optional[bool] = None,
-
band_as_variable: bool = False,
-
**open_kwargs,
-
) -> Union[Dataset, DataArray, list[Dataset]]:
-
# pylint: disable=too-many-statements,too-many-locals,too-many-branches
-
"""Open a file with rasterio (experimental).
-
-
This should work with any file that rasterio can open (most often:
-
geoTIFF). The x and y coordinates are generated automatically from the
-
file's geoinformation, shifted to the center of each pixel (see
-
`"PixelIsArea" Raster Space
-
<http://web.archive.org/web/20160326194152/http://remotesensing.org/geotiff/spec/geotiff2.5.html#2.5.2>`_
-
for more information).
-
-
.. versionadded:: 0.13 band_as_variable
-
-
Parameters
-
----------
-
filename: str, rasterio.io.DatasetReader, or rasterio.vrt.WarpedVRT
-
Path to the file to open. Or already open rasterio dataset.
-
parse_coordinates: bool, optional
-
Whether to parse the x and y coordinates out of the file's
-
``transform`` attribute or not. The default is to automatically
-
parse the coordinates only if they are rectilinear (1D).
-
It can be useful to set ``parse_coordinates=False``
-
if your files are very large or if you don't need the coordinates.
-
chunks: int, tuple or dict, optional
-
Chunk sizes along each dimension, e.g., ``5``, ``(5, 5)`` or
-
``{'x': 5, 'y': 5}``. If chunks is provided, it used to load the new
-
DataArray into a dask array. Chunks can also be set to
-
``True`` or ``"auto"`` to choose sensible chunk sizes according to
-
``dask.config.get("array.chunk-size")``.
-
cache: bool, optional
-
If True, cache data loaded from the underlying datastore in memory as
-
NumPy arrays when accessed to avoid reading from the underlying data-
-
store multiple times. Defaults to True unless you specify the `chunks`
-
argument to use dask, in which case it defaults to False.
-
lock: bool or dask.utils.SerializableLock, optional
-
-
If chunks is provided, this argument is used to ensure that only one
-
thread per process is reading from a rasterio file object at a time.
-
-
By default and when a lock instance is provided,
-
a :class:`xarray.backends.CachingFileManager` is used to cache File objects.
-
Since rasterio also caches some data, this will make repeated reads from the
-
same object fast.
-
-
When ``lock=False``, no lock is used, allowing for completely parallel reads
-
from multiple threads or processes. However, a new file handle is opened on
-
each request.
-
-
masked: bool, optional
-
If True, read the mask and set values to NaN. Defaults to False.
-
mask_and_scale: bool, default=False
-
Lazily scale (using the `scales` and `offsets` from rasterio) and mask.
-
If the _Unsigned attribute is present treat integer arrays as unsigned.
-
variable: str or list or tuple, optional
-
Variable name or names to use to filter loading.
-
group: str or list or tuple, optional
-
Group name or names to use to filter loading.
-
default_name: str, optional
-
The name of the data array if none exists. Default is None.
-
decode_times: bool, default=True
-
If True, decode times encoded in the standard NetCDF datetime format
-
into datetime objects. Otherwise, leave them encoded as numbers.
-
decode_timedelta: bool, optional
-
If True, decode variables and coordinates with time units in
-
{“days”, “hours”, “minutes”, “seconds”, “milliseconds”, “microseconds”}
-
into timedelta objects. If False, leave them encoded as numbers.
-
If None (default), assume the same value of decode_time.
-
band_as_variable: bool, default=False
-
If True, will load bands in a raster to separate variables.
-
**open_kwargs: kwargs, optional
-
Optional keyword arguments to pass into :func:`rasterio.open`.
-
-
Returns
-
-------
-
:obj:`xarray.Dataset` | :obj:`xarray.DataArray` | list[:obj:`xarray.Dataset`]:
-
The newly created dataset(s).
-
"""
-
parse_coordinates = True if parse_coordinates is None else parse_coordinates
-
masked = masked or mask_and_scale
-
vrt_params = None
-
file_opener = rasterio.open
-
if isinstance(filename, SingleBandDatasetReader):
-
file_opener = functools.partial(
-
_single_band_open,
-
bidx=filename._bidx,
-
)
-
vrt_params = filename._vrt_params
-
if isinstance(filename, (rasterio.io.DatasetReader, SingleBandDatasetReader)):
-
filename = filename.name
-
elif isinstance(filename, rasterio.vrt.WarpedVRT):
-
vrt = filename
-
filename = vrt.src_dataset.name
-
vrt_params = {
-
"src_crs": vrt.src_crs.to_string() if vrt.src_crs else None,
-
"crs": vrt.dst_crs.to_string() if vrt.dst_crs else None,
-
"resampling": vrt.resampling,
-
"tolerance": vrt.tolerance,
-
"src_nodata": vrt.src_nodata,
-
"nodata": vrt.dst_nodata,
-
"width": vrt.dst_width,
-
"height": vrt.dst_height,
-
"src_transform": vrt.src_transform,
-
"transform": vrt.dst_transform,
-
"dtype": vrt.working_dtype,
-
**vrt.warp_extras,
-
}
-
-
if lock in (True, None):
-
lock = RASTERIO_LOCK
-
elif lock is False:
-
lock = NO_LOCK
-
-
# ensure default for sharing is False
-
# ref https://github.com/mapbox/rasterio/issues/1504
-
open_kwargs["sharing"] = open_kwargs.get("sharing", False)
-
-
with warnings.catch_warnings(record=True) as rio_warnings:
-
if lock is not NO_LOCK and isinstance(filename, (str, os.PathLike)):
-
manager: FileManager = CachingFileManager(
-
file_opener, filename, lock=lock, mode="r", kwargs=open_kwargs
-
)
-
else:
-
manager = URIManager(file_opener, filename, mode="r", kwargs=open_kwargs)
-
riods = manager.acquire()
-
captured_warnings = rio_warnings.copy()
-
-
# raise the NotGeoreferencedWarning if applicable
-
for rio_warning in captured_warnings:
-
if not riods.subdatasets or not isinstance(
-
rio_warning.message, NotGeoreferencedWarning
-
):
-
warnings.warn(str(rio_warning.message), type(rio_warning.message)) # type: ignore
-
-
# open the subdatasets if they exist
-
if riods.subdatasets:
-
subdataset_result = _load_subdatasets(
-
riods=riods,
-
group=group,
-
variable=variable,
-
parse_coordinates=parse_coordinates,
-
chunks=chunks,
-
cache=cache,
-
lock=lock,
-
masked=masked,
-
mask_and_scale=mask_and_scale,
-
decode_times=decode_times,
-
decode_timedelta=decode_timedelta,
-
**open_kwargs,
-
)
-
manager.close()
-
return subdataset_result
-
-
if band_as_variable:
-
dataset_result = _load_bands_as_variables(
-
riods=riods,
-
parse_coordinates=parse_coordinates,
-
chunks=chunks,
-
cache=cache,
-
lock=lock,
-
masked=masked,
-
mask_and_scale=mask_and_scale,
-
decode_times=decode_times,
-
decode_timedelta=decode_timedelta,
-
vrt_params=vrt_params,
-
**open_kwargs,
-
)
-
manager.close()
-
return dataset_result
-
-
if cache is None:
-
cache = chunks is None
-
-
riods = _ensure_warped_vrt(riods, vrt_params)
-
-
# Get bands
-
if riods.count < 1:
-
raise ValueError("Unknown dims")
-
-
# parse tags & load alternate coords
-
attrs = _get_rasterio_attrs(riods=riods)
-
coords = _load_netcdf_1d_coords(attrs)
-
_parse_driver_tags(riods=riods, attrs=attrs, coords=coords)
-
for coord in coords:
-
if f"NETCDF_DIM_{coord}" in attrs:
-
coord_name = coord
-
attrs.pop(f"NETCDF_DIM_{coord}")
-
break
-
if f"NETCDF_DIM_{coord}_VALUES" in attrs:
-
coord_name = coord
-
attrs.pop(f"NETCDF_DIM_{coord}_VALUES")
-
attrs.pop(f"NETCDF_DIM_{coord}_DEF", None)
-
attrs.pop("NETCDF_DIM_EXTRA", None)
-
break
-
else:
-
coord_name = "band"
-
coords[coord_name] = numpy.asarray(riods.indexes)
-
-
has_gcps = riods.gcps[0]
-
if has_gcps:
-
parse_coordinates = False
-
-
# Get geospatial coordinates
-
if parse_coordinates:
-
coords.update(
-
_generate_spatial_coords(riods.transform, riods.width, riods.height)
-
)
-
-
unsigned = None
-
encoding: dict[Hashable, Any] = {}
-
if mask_and_scale and "_Unsigned" in attrs:
-
unsigned = variables.pop_to(attrs, encoding, "_Unsigned") == "true"
-
-
if masked:
-
encoding["dtype"] = str(_rasterio_to_numpy_dtype(riods.dtypes))
-
-
da_name = attrs.pop("NETCDF_VARNAME", default_name)
-
data: Any = indexing.LazilyOuterIndexedArray(
-
RasterioArrayWrapper(
-
manager,
-
lock,
-
name=da_name,
-
vrt_params=vrt_params,
-
masked=masked,
-
mask_and_scale=mask_and_scale,
-
unsigned=unsigned,
-
)
-
)
-
-
# this lets you write arrays loaded with rasterio
-
data = indexing.CopyOnWriteArray(data)
-
if cache and chunks is None:
-
data = indexing.MemoryCachedArray(data)
-
-
result = DataArray(
-
data=data, dims=(coord_name, "y", "x"), coords=coords, attrs=attrs, name=da_name
-
)
-
result.encoding = encoding
-
-
# update attributes from NetCDF attributes
-
_load_netcdf_attrs(riods.tags(), result)
-
result = _decode_datetime_cf(
-
result, decode_times=decode_times, decode_timedelta=decode_timedelta
-
)
-
-
# make sure the _FillValue is correct dtype
-
if "_FillValue" in result.attrs:
-
result.attrs["_FillValue"] = result.dtype.type(result.attrs["_FillValue"])
-
-
# handle encoding
-
_handle_encoding(result, mask_and_scale, masked, da_name, unsigned=unsigned)
-
# Affine transformation matrix (always available)
-
# This describes coefficients mapping pixel coordinates to CRS
-
# For serialization store as tuple of 6 floats, the last row being
-
# always (0, 0, 1) per definition (see
-
# https://github.com/sgillies/affine)
-
result.rio.write_transform(riods.transform, inplace=True)
-
rio_crs = riods.crs or result.rio.crs
-
if rio_crs:
-
result.rio.write_crs(rio_crs, inplace=True)
-
if has_gcps:
-
result.rio.write_gcps(*riods.gcps, inplace=True)
-
-
if chunks is not None:
-
result = _prepare_dask(result, riods, filename, chunks)
-
else:
-
result.encoding["preferred_chunks"] = {
-
result.rio.y_dim: riods.block_shapes[0][0],
-
result.rio.x_dim: riods.block_shapes[0][1],
-
coord_name: 1,
-
}
-
-
# add file path to encoding
-
result.encoding["source"] = riods.name
-
result.encoding["rasterio_dtype"] = str(riods.dtypes[0])
-
# remove duplicate coordinate information
-
for coord in result.coords:
-
result.attrs = {
-
attr: value
-
for attr, value in result.attrs.items()
-
if not attr.startswith(f"{coord}#")
-
}
-
# remove duplicate tags
-
if result.name:
-
result.attrs = {
-
attr: value
-
for attr, value in result.attrs.items()
-
if not attr.startswith(f"{result.name}#")
-
}
-
# Make the file closeable
-
result.set_close(manager.close)
-
result.rio._manager = manager
-
return result
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
