import numpy as np
from buzzard._a_stored_raster import AStoredRaster, ABackStoredRaster
class NumpyRaster(AStoredRaster):
"""Concrete class defining the behavior of a wrapped numpy array
>>> help(Dataset.wrap_numpy_raster)
Features Defined
----------------
- Has an `array` property that points to the numpy array provided at construction.
"""
def __init__(self, ds, fp, array, channels_schema, wkt, mode):
self._arr_shape = array.shape
self._arr_address = array.__array_interface__['data'][0]
back = BackNumpyRaster(
ds._back, fp, array, channels_schema, wkt, mode
)
super(NumpyRaster, self).__init__(ds=ds, back=back)
@property
def array(self):
"""Returns the Raster's full input data as a Numpy array"""
assert (
self._arr_address == self._back._arr.__array_interface__['data'][0]
)
return self._back._arr.reshape(*self._arr_shape)
class BackNumpyRaster(ABackStoredRaster):
"""Implementation of NumpyRaster"""
def __init__(self, back_ds, fp, array, channels_schema, wkt, mode):
array = np.atleast_3d(array)
self._arr = array
channel_count = array.shape[-1]
if 'nodata' not in channels_schema:
channels_schema['nodata'] = [None] * channel_count
if 'interpretation' not in channels_schema:
channels_schema['interpretation'] = ['undefined'] * channel_count
if 'offset' not in channels_schema:
channels_schema['offset'] = [0.] * channel_count
if 'scale' not in channels_schema:
channels_schema['scale'] = [1.] * channel_count
if 'mask' not in channels_schema:
channels_schema['mask'] = ['all_valid']
super(BackNumpyRaster, self).__init__(
back_ds=back_ds,
wkt_stored=wkt,
channels_schema=channels_schema,
dtype=array.dtype,
fp_stored=fp,
mode=mode,
)
self._should_tranform = (
any(v != 0 for v in channels_schema['offset']) or
any(v != 1 for v in channels_schema['scale'])
)
def get_data(self, fp, channel_ids, dst_nodata, interpolation):
samplefp = self.build_sampling_footprint(fp, interpolation)
if samplefp is None:
return np.full(
np.r_[fp.shape, len(channel_ids)],
dst_nodata,
self.dtype
)
chans_indexer = self._best_indexers_of_channel_ids(channel_ids)
key = list(samplefp.slice_in(self.fp)) + [chans_indexer]
key = tuple(key)
array = self._arr[key]
if self._should_tranform:
array = (
array *
np.asarray(self.channels_schema['scale'])[chans_indexer] +
np.asarray(self.channels_schema['offset'])[chans_indexer]
)
array = self.remap(
samplefp,
fp,
array=array,
mask=None,
src_nodata=self.nodata,
dst_nodata=dst_nodata,
mask_mode='erode',
interpolation=interpolation,
)
array = array.astype(self.dtype, copy=False)
return array
def set_data(self, array, fp, channel_ids, interpolation, mask):
if not fp.share_area(self.fp):
return
if not fp.same_grid(self.fp) and mask is None:
mask = np.ones(fp.shape, bool)
dstfp = self.fp.intersection(fp)
# Remap ****************************************************************
ret = self.remap(
fp,
dstfp,
array=array,
mask=mask,
src_nodata=self.nodata,
dst_nodata=self.nodata or 0,
mask_mode='erode',
interpolation=interpolation,
)
if mask is not None:
array, mask = ret
else:
array = ret
del ret
array = array.astype(self.dtype, copy=False)
fp = dstfp
del dstfp
# Write ****************************************************************
slices = fp.slice_in(self.fp)
for i in channel_ids:
if mask is not None:
self._arr[slices + (i,)][mask] = array[..., i][mask]
else:
self._arr[slices + (i,)] = array[..., i]
def fill(self, value, channel_ids):
for i in channel_ids:
self._arr[..., i] = value
def close(self):
super(BackNumpyRaster, self).close()
del self._arr
@staticmethod
def _best_indexers_of_channel_ids(channel_ids):
"""Create an object to pick the channels of the numpy array. Returns either a slice
object or a list of int to perform fancy-indexing"""
l = list(channel_ids)
if np.all(np.diff(l) == 1):
start, stop = l[0], l[-1] + 1
l = slice(start, stop)
elif np.all(np.diff(l) == -1):
start, stop = l[0], l[-1] - 1
if stop < 0:
stop = None
l = slice(start, stop, -1)
return l
