def test_aux_1(lt_ctx):
data = _mk_random(size=(16, 16, 16, 16), dtype="float32")
aux_data = BufferWrapper(kind="nav", dtype="float32", extra_shape=(2, ))
aux_data.set_buffer(_mk_random(size=(16, 16, 2), dtype="float32"))
dataset = MemoryDataSet(data=data,
tileshape=(7, 16, 16),
num_partitions=2,
sig_dims=2)
echo_udf = EchoUDF(aux=aux_data)
res = lt_ctx.run_udf(dataset=dataset, udf=echo_udf)
assert 'echo' in res
print(data.shape, res['echo'].data.shape)
assert np.allclose(res['echo'].raw_data, aux_data.raw_data)
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 = BufferWrapper(kind, extra_shape, dtype)
buf.set_buffer(data)
return buf
