def test_new_for_partition():
auxdata = _mk_random(size=(16, 16), dtype="float32")
buf = AuxBufferWrapper(kind="nav", dtype="float32")
buf.set_buffer(auxdata)
dataset = MemoryDataSet(data=_mk_random(size=(16, 16, 16, 16),
dtype="float32"),
tileshape=(7, 16, 16),
num_partitions=2,
sig_dims=2)
assert auxdata.shape == tuple(dataset.shape.nav)
roi = _mk_random(size=dataset.shape.nav, dtype="bool")
for idx, partition in enumerate(dataset.get_partitions()):
print("partition number", idx)
new_buf = buf.new_for_partition(partition, roi=roi)
ps = partition.slice.get(nav_only=True)
roi_part = roi.reshape(-1)[ps]
assert np.product(new_buf._data.shape) == roi_part.sum()
# old buffer stays the same:
assert np.allclose(buf._data, auxdata.reshape(-1))
assert buf._data_coords_global
assert not new_buf._data_coords_global
# new buffer is sliced to partition and has ROI applied:
assert new_buf._data.shape[0] <= buf._data.shape[0]
assert new_buf._data.shape[0] <= partition.shape[0]
# let's try and manually apply the ROI to `auxdata`:
assert np.allclose(new_buf._data, auxdata.reshape(-1)[ps][roi_part])
def aux_data(cls, data, kind, extra_shape=(), dtype="float32"):
"""
Use this method to create auxiliary data. Auxiliary data should
have a shape like `(dataset.shape.nav, extra_shape)` and on access,
an appropriate view will be created. For example, if you access
aux data in `process_frame`, you will get the auxiliary data for
the current frame you are processing.
Example
-------
We create a UDF to demonstrate the behavior:
>>> class MyUDF(UDF):
... def get_result_buffers(self):
... # Result buffer for debug output
... return {'aux_dump': self.buffer(kind='nav', dtype='object')}
...
... def process_frame(self, frame):
... # Extract value of aux data for demonstration
... self.results.aux_dump[:] = str(self.params.aux_data[:])
...
>>> # for each frame, provide three values from a sequential series:
>>> aux1 = MyUDF.aux_data(
... data=np.arange(np.prod(dataset.shape.nav) * 3, dtype=np.float32),
... kind="nav", extra_shape=(3,), dtype="float32"
... )
>>> udf = MyUDF(aux_data=aux1)
>>> res = ctx.run_udf(dataset=dataset, udf=udf)
process_frame for frame (0, 7) received a view of aux_data with values [21., 22., 23.]:
>>> res['aux_dump'].data[0, 7]
'[21. 22. 23.]'
"""
buf = AuxBufferWrapper(kind, extra_shape, dtype)
buf.set_buffer(data)
return buf