def main():
MAX_CHUNKS = 4
place_callback = scipy.LowLevelCallable(chunk_place.ctypes,
signature='void (void *, size_t)')
chunk_config = spead2.recv.ChunkStreamConfig(
items=[spead2.HEAP_CNT_ID, spead2.HEAP_LENGTH_ID],
max_chunks=MAX_CHUNKS,
place=place_callback)
data_ring = spead2.recv.ChunkRingbuffer(MAX_CHUNKS)
free_ring = spead2.recv.ChunkRingbuffer(MAX_CHUNKS)
stream = spead2.recv.ChunkRingStream(spead2.ThreadPool(),
spead2.recv.StreamConfig(),
chunk_config, data_ring, free_ring)
for i in range(MAX_CHUNKS):
chunk = spead2.recv.Chunk(present=np.empty(HEAPS_PER_CHUNK, np.uint8),
data=np.empty(CHUNK_PAYLOAD_SIZE, np.uint8))
stream.add_free_chunk(chunk)
stream.add_udp_reader(8888,
buffer_size=1024 * 1024,
bind_hostname='127.0.0.1')
for chunk in data_ring:
n_present = np.sum(chunk.present)
print(f"Received chunk {chunk.chunk_id} with "
f"{n_present} / {HEAPS_PER_CHUNK} heaps")
stream.add_free_chunk(chunk)
def test_set_place_bad_signature(self):
place = scipy.LowLevelCallable(place_plain.ctypes,
signature='void (void)')
# One might expect TypeError, but ValueError is what scipy uses for
# invalid signatures.
with pytest.raises(ValueError):
recv.ChunkStreamConfig(place=place)
def test_packet_presence(self, data_ring, queue):
"""Test packet presence feature."""
# Each heap is split into two packets. Create a free ring where the
# chunks have space for this.
free_ring = spead2.recv.ChunkRingbuffer(4)
while not free_ring.full():
free_ring.put(
recv.Chunk(present=np.zeros(HEAPS_PER_CHUNK * PACKETS_PER_HEAP,
np.uint8),
data=np.zeros(CHUNK_PAYLOAD_SIZE, np.uint8)))
stream_config = spead2.recv.StreamConfig(max_heaps=128,
allow_out_of_order=True)
# Note: user_data is deliberately not assigned to a local variable, so
# that reference-counting errors are more likely to be detected.
user_data = np.zeros(1, dtype=user_data_type.dtype)
user_data["scale"] = PACKETS_PER_HEAP
user_data["placed_heaps_index"] = stream_config.add_stat(
"placed_heaps")
place_bind_llc = scipy.LowLevelCallable(
place_bind.ctypes,
user_data=user_data.ctypes.data_as(ctypes.c_void_p),
signature='void (void *, size_t, void *)')
stream = spead2.recv.ChunkRingStream(
spead2.ThreadPool(), stream_config,
spead2.recv.ChunkStreamConfig(
items=[
0x1000, spead2.HEAP_LENGTH_ID, spead2.PAYLOAD_LENGTH_ID
],
max_chunks=4,
place=place_bind_llc,
).enable_packet_presence(PACKET_SIZE), data_ring, free_ring)
stream.add_inproc_reader(queue)
queue.add_packet(self.make_packet(4, 0, PACKET_SIZE))
queue.add_packet(self.make_packet(4, PACKET_SIZE, 2 * PACKET_SIZE))
queue.add_packet(self.make_packet(17, PACKET_SIZE, 2 * PACKET_SIZE))
queue.stop()
chunks = list(stream.data_ringbuffer)
assert len(chunks) == 2
self.check_chunk_packets(
chunks[0], 0,
np.array(
[0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
np.uint8))
self.check_chunk_packets(
chunks[1], 1,
np.array(
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0],
np.uint8))
assert stream.stats["placed_heaps"] == 2
class ScipyModule:
"""Converts c_capsules into a set of scipy.LowLevelCallable"""
n = c_capsules.n
f = scipy.LowLevelCallable(c_capsules.f)
grad_f = scipy.LowLevelCallable(c_capsules.grad_f)
g = scipy.LowLevelCallable(c_capsules.g)
jac_g = scipy.LowLevelCallable(c_capsules.jac_g)
h = scipy.LowLevelCallable(c_capsules.h)
intermediate_callback = scipy.LowLevelCallable(
c_capsules.intermediate_callback)
def make_integrand():
global _integrand
import numba
import numba.types as nt
import scipy
@numba.cfunc(nt.double(nt.intc, nt.CPointer(nt.double)),
nopython=True, nogil=True)
def _fn(n, xx):
if n != 4:
return np.nan
y = xx[0]
k = xx[1]
sigma2 = xx[2]
z = xx[3]
return y ** (k/2 - 1) * np.exp(-(k * y + (z-y)**2 / sigma2) / 2)
_integrand = scipy.LowLevelCallable(_fn.ctypes)
def A_x(xi, xi2, h_nodes, num_nodes):
# From fn_Qx_matrix
# Notes: the longest in this function is the numerical integration with scipy.integrate.quad.
# To speed it up, the integrand is compiled with numba. The quad function requires
# a function f8(f8, voidptr) so the arguments xi, xi2, h_nodes and z, needed for the
# calculation of the integrand, are packed.
# integral around branch point
z_1 = xi2 * h_nodes * np.arange(num_nodes)
# Pack arguments for LowLevelCallable.
# data is updated at every loop. The main loop is NOT thread-safe. If the main loop
# becomes parallel some day, make "data" local.
data = np.array([xi, xi2, h_nodes, 0.0])
data_ptr = ctypes.cast(data.ctypes, ctypes.c_void_p)
quad_args = dict(limit=200)
# For num_nodes = 4, I_12 looks like [a3, a2, a1, a0, a1, a2, a3] (size: 2*num_nodes-1)
# Build the second half first, then copy it to the first half
I_12 = np.zeros(2 * num_nodes - 1, np.complex)
for i, z in enumerate(z_1):
data[3] = z
if i < 2: # two first iterations, coefficients a0 and a1
int_F1_real = scipy.LowLevelCallable(A_x_F1_real.ctypes, data_ptr)
int_F1_imag = scipy.LowLevelCallable(A_x_F1_imag.ctypes, data_ptr)
int_F2_real = scipy.LowLevelCallable(A_x_F2_real.ctypes, data_ptr)
int_F2_imag = scipy.LowLevelCallable(A_x_F2_imag.ctypes, data_ptr)
I_12[i + num_nodes - 1] = 4j * (
si.quad(int_F1_real, 0, 1, **quad_args)[0] +
1j * si.quad(int_F1_imag, 0, 1, **quad_args)[0]) + 4 * (
si.quad(int_F2_real, 0, 50, **quad_args)[0] +
1j * si.quad(int_F2_imag, 0, 50, **quad_args)[0])
else:
int_F_real = scipy.LowLevelCallable(A_x_F_real.ctypes, data_ptr)
int_F_imag = scipy.LowLevelCallable(A_x_F_imag.ctypes, data_ptr)
I_12[i + num_nodes -
1] = 4j * (si.quad(int_F_real, 0, 70, **quad_args)[0] +
1j * si.quad(int_F_imag, 0, 70, **quad_args)[0])
I_12[:num_nodes - 1] = I_12[:num_nodes - 1:-1]
v_ind = np.arange(num_nodes)
m_ind = (np.full(
(num_nodes, num_nodes), num_nodes - 1) + v_ind[:, np.newaxis] - v_ind)
return I_12[m_ind]
def A_z(xi, xi2, h_nodes, num_nodes):
# from fn_Qz_matrix
# integral around branch point
z_1 = xi2 * h_nodes * np.arange(num_nodes)
# Pack arguments for LowLevelCallable.
# data is updated at every loop. The main loop is NOT thread-safe. If the main loop
# becomes parallel some day, make "data" local.
data = np.array([xi, xi2, h_nodes, 0.0])
data_ptr = ctypes.cast(data.ctypes, ctypes.c_void_p)
quad_args = dict(limit=200)
# For num_nodes = 4, I_12 looks like [a3, a2, a1, a0, a1, a2, a3] (size: 2*num_nodes-1)
# Build the second half first, then copy it to the first half
I_12 = np.zeros(2 * num_nodes - 1, np.complex)
for i, z in enumerate(z_1):
data[3] = z
if i < 2: # two first iterations, coefficients a0 and a1
int_F1_real = scipy.LowLevelCallable(A_z_F1_real.ctypes, data_ptr)
int_F1_imag = scipy.LowLevelCallable(A_z_F1_imag.ctypes, data_ptr)
int_F2_real = scipy.LowLevelCallable(A_z_F2_real.ctypes, data_ptr)
int_F2_imag = scipy.LowLevelCallable(A_z_F2_imag.ctypes, data_ptr)
I_12[i + num_nodes - 1] = 4j * (
si.quad(int_F1_real, 0, sqrt(xi), **quad_args)[0] +
1j * si.quad(int_F1_imag, 0, sqrt(xi), **quad_args)[0]) + 4 * (
si.quad(int_F2_real, 0, 50, **quad_args)[0] +
1j * si.quad(int_F2_imag, 0, 50, **quad_args)[0])
else:
int_F_real = scipy.LowLevelCallable(A_z_F_real.ctypes, data_ptr)
int_F_imag = scipy.LowLevelCallable(A_z_F_imag.ctypes, data_ptr)
I_12[i + num_nodes -
1] = 4j * (si.quad(int_F_real, 0, 70, **quad_args)[0] +
1j * si.quad(int_F_imag, 0, 70, **quad_args)[0])
I_12[:num_nodes - 1] = I_12[:num_nodes - 1:-1]
v_ind = np.arange(num_nodes)
m_ind = (np.full(
(num_nodes, num_nodes), num_nodes - 1) + v_ind[:, np.newaxis] - v_ind)
return I_12[m_ind]
packet_size = items[2]
user_data = numba.carray(user_data_ptr, 1)
if payload_size == HEAP_PAYLOAD_SIZE and packet_size == PACKET_SIZE:
data[0].chunk_id = heap_cnt // HEAPS_PER_CHUNK
heap_index = heap_cnt % HEAPS_PER_CHUNK
data[0].heap_index = heap_index * user_data[0].scale
data[0].heap_offset = heap_index * HEAP_PAYLOAD_SIZE
batch_stats = numba.carray(intp_to_voidptr(data[0].batch_stats),
user_data[0].placed_heaps_index + 1,
dtype=np.uint64)
batch_stats[user_data[0].placed_heaps_index] += 1
# ctypes doesn't distinguish equivalent integer types, so we have to
# specify the signature explicitly.
place_plain_llc = scipy.LowLevelCallable(place_plain.ctypes,
signature='void (void *, size_t)')
place_bind_llc = scipy.LowLevelCallable(
place_bind.ctypes, signature='void (void *, size_t, void *)')
class TestChunkStreamConfig:
def test_default_construct(self):
config = recv.ChunkStreamConfig()
assert config.items == []
assert config.max_chunks == config.DEFAULT_MAX_CHUNKS
assert config.place is None
assert config.packet_presence_payload_size == 0
def test_zero_max_chunks(self):
config = recv.ChunkStreamConfig()
with pytest.raises(ValueError):