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
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.utils import is_scalar
from rioxarray.exceptions import RioXarrayError
from rioxarray.rioxarray import affine_to_coords
# TODO: should this be GDAL_LOCK instead?
RASTERIO_LOCK = SerializableLock()
class RasterioArrayWrapper(BackendArray):
"""A wrapper around rasterio dataset objects"""
def __init__(
self,
manager,
lock,
name,
vrt_params=None,
masked=False,
mask_and_scale=False,
unsigned=False,
):
from rasterio.vrt import WarpedVRT
# no locking for xarray <= 0.11
def ensure_lock(lock):
return lock
class SerializableLock(object):
def __enter__(self):
pass
def __exit__(self, exc_type, exc_val, exc_tb):
pass
# FIXME: Add a dedicated lock, even if ecCodes is supposed to be thread-safe
# in most circumstances. See:
# https://confluence.ecmwf.int/display/ECC/Frequently+Asked+Questions
ECCODES_LOCK = SerializableLock()
class CfGribArrayWrapper(BackendArray):
def __init__(self, datastore, array):
self.datastore = datastore
self.shape = array.shape
self.dtype = array.dtype
self.array = array
def __getitem__(self, key):
return indexing.explicit_indexing_adapter(
key, self.shape, indexing.IndexingSupport.OUTER, self._getitem)
def _getitem(self, key):
with self.datastore.lock:
import pandas as pd
import xarray as xr
from xarray.backends.common import (
BACKEND_ENTRYPOINTS,
AbstractDataStore,
BackendArray,
BackendEntrypoint,
)
from xarray.core import indexing
from xarray.backends.locks import SerializableLock
import pytdlpack
import TdlpackIO
logger = logging.getLogger(__name__)
LOCK = SerializableLock()
class TdlpackBackendEntrypoint(BackendEntrypoint):
''' xarray backend engine entrypoint for opening and decoding sequential tdlpack files.
.. warning::
This backend is pre-release and its signature may change without backward comaptability.
Parameters
__________
filename: str, Path, file-like
sequential tdlpack file to be opened
name_scheme: list of strings with tdlpack metadata, optional
dictates name of variables and shape of arrays;
from . import UNDEFINED, _Variable, WgribError
from .wgrib2 import MemoryBuffer, wgrib, free_files
from .inventory import MetaData
from .template import Template
logger = logging.getLogger(__name__)
# wgrib2 returns C float arrays
DTYPE = np.dtype("float32")
HeaderIndices = Tuple[int, ...]
FileIndex = DefaultDict[str, Dict[HeaderIndices, str]] # file -> Dict
FileIndices = DefaultDict[str, FileIndex] # variable name -> FileIndex
WGRIB2_LOCK = SerializableLock()
class Dataset(NamedTuple):
dims: Dict[str, int]
vars: Dict[str, _Variable]
attrs: Dict[str, Any]
def find_1st(array, value):
return np.nonzero(array == value)[0][0]
def build_file_index(
items: Sequence[MetaData],
template: Template,