diff --git a/ch_vdif_assembler.py b/ch_vdif_assembler.py index 4287f02..fa8aa2c 100644 --- a/ch_vdif_assembler.py +++ b/ch_vdif_assembler.py @@ -48,10 +48,10 @@ class constants: num_disks = ch_vdif_assembler_cython.num_disks # 10 (moose) -class assembler: +class assembler(object): def __init__(self, write_to_disk=False, rbuf_size=constants.num_disks, abuf_size=4, assembler_nt=65536): self._assembler = ch_vdif_assembler_cython.assembler(write_to_disk, rbuf_size, abuf_size, assembler_nt) - self.python_processor = None + self.python_processors = [] def register_processor(self, p): @@ -59,31 +59,42 @@ def register_processor(self, p): self._assembler.register_cpp_processor(p) # register C++ processor (this actually spawns a processing thread) elif not isinstance(p, processor): raise RuntimeError('Argument to assembler.register_processor() must be either an object of class ch_vdif_assembler.processor, or a C++ processor (e.g. returned by make_waterfall_plotter)') - elif self.python_processor is not None: - raise RuntimeError('Currently, ch_vdif_assembler only allows registering one python processor (but an arbitrary number of C++ processors)') - else: - self.python_processor = p + self.python_processors.append(p) def run(self, stream): - if self.python_processor is None: + if not self.python_processors: self._assembler.start_async(stream) self._assembler.wait_until_end() return self._assembler.register_python_processor() + processors_byte_data = [p.byte_data for p in self.python_processors] + need_byte_data = True in processors_byte_data + need_complex_data = False in processors_byte_data + try: self._assembler.start_async(stream) - + while True: chunk = self._assembler.get_next_python_chunk() if chunk is None: break - (t0, nt, efield, mask) = chunk.get_data() - self.python_processor.process_chunk(t0, nt, efield, mask) - - self.python_processor.finalize() + if need_byte_data: + byte_data = chunk.get_byte_data() + if need_complex_data: + complex_data = chunk.get_data() + for p in self.python_processors: + if p.byte_data: + t0, nt, efield = byte_data + mask = None + else: + t0, nt, efield, mask = complex_data + p.process_chunk(t0, nt, efield, mask) + + for p in self.python_processors: + p.finalize() finally: self._assembler.unregister_python_processor() @@ -142,7 +153,7 @@ def cpp_waterfall_plotter(outdir, is_critical=False): return ch_vdif_assembler_cython.cpp_waterfall_plotter(outdir, is_critical) -class processor: +class processor(object): """ To define a python processor, you subclass this base class. @@ -158,17 +169,25 @@ class processor: interruption in data stream, then a timestamp gap will appear. The 'efield' arg is a shape (nfreq,2,nt) complex array with electric field values, where - the middle index is polarziation. Missing data is represented by (0+0j). The 'mask' arg + the middle index is polarization. Missing data is represented by (0+0j). The 'mask' arg is a shape (nfreq,2,nt) integer array which is 0 for missing data, and 1 for non-missing. - WARNING 2: Handling missing data is an important aspect of the vdif_processor since it - happens all the time. If a GPU correlator node is down, which is a frequent occurrence, - then some frequencies will be "all missing". There are also routine packet loss events - on second-timescales which result in some high-speed samples being flagged as missing data. - """ + For subclasses with attribute `byte_data = True`, 'efield' is a shape + (nfreq,2,nt) byte array, as in the C++ versions. 'mask' is None. + + WARNING 2: Handling missing data is an important aspect of the vdif_processor since it + happens all the time. If a GPU correlator node is down, which is a frequent occurrence, + then some frequencies will be "all missing". There are also routine packet loss events + on second-timescales which result in some high-speed samples being flagged as missing data. + """ + + byte_data = False def process_chunk(self, t0, nt, efield, mask): - print 'process_chunk called! t0=%s nt=%s efield (%s,%s) mask (%s,%s)' % (t0, nt, efield.dtype, efield.shape, mask.dtype, mask.shape) + if mask is None: + print 'process_chunk called! t0=%s nt=%s efield (%s,%s)' % (t0, nt, efield.dtype, efield.shape) + else: + print 'process_chunk called! t0=%s nt=%s efield (%s,%s) mask (%s,%s)' % (t0, nt, efield.dtype, efield.shape, mask.dtype, mask.shape) def finalize(self): pass @@ -181,7 +200,7 @@ def finalize(self): # See also the script show-moose-acquisitions.py -class moose_inventory: +class moose_inventory(object): def __init__(self): self.topdirs = [ ('/drives/E/%d' % i) for i in xrange(10) ] self.subdirs = set() diff --git a/ch_vdif_assembler_cython.hpp b/ch_vdif_assembler_cython.hpp index eb58789..62aea38 100644 --- a/ch_vdif_assembler_cython.hpp +++ b/ch_vdif_assembler_cython.hpp @@ -89,6 +89,16 @@ struct cython_assembled_chunk { std::complex *efield = reinterpret_cast *> (efield_hack); p->fill_efield_array_reference(efield, mask); } + + inline uint8_t* get_buf(void) + { + // This const_cast is super sketchy but I can't figure out a way to get around it since I want + // to use this buffer as a numpy array (in the python assembled_chunk.get_byte_data). That array + // is set as read only. It would be marginally better if the const_cast could be done later in cython, + // but I can't figure out how to do that. -KM + return const_cast (p->buf); + } + }; diff --git a/ch_vdif_assembler_cython.pyx b/ch_vdif_assembler_cython.pyx index 49c2c02..3ca5711 100644 --- a/ch_vdif_assembler_cython.pyx +++ b/ch_vdif_assembler_cython.pyx @@ -1,10 +1,14 @@ -from libc.stdint cimport int32_t, int64_t +from libc.stdint cimport int32_t, int64_t, uint8_t +from libc.stdlib cimport free, malloc from libcpp.vector cimport vector from libcpp cimport bool import numpy as np cimport numpy as np +# Oddly, things compile without this, but segfaults. -KM +np.import_array() # Exposes the numpy c-api. + cimport ch_vdif_assembler_pxd @@ -95,6 +99,48 @@ cdef class assembled_chunk: self.p[0].fill_efield(&efield[0,0,0], &mask[0,0,0]) return (t0, nt, efield, mask) + def get_byte_data(self): + if self.p == NULL: + return None + + t0 = self.p[0].t0 + nt = self.p[0].nt + + shape = (chime_nfreq, 2, nt) + # Get all the dimensions in the right format. + cdef int nd = len(shape) + # Very important to use this type. + cdef np.npy_intp *shape_c = malloc(nd * sizeof(np.npy_intp)) + cdef int64_t size = 1 + for ii, s in enumerate(shape): + shape_c[ii] = s + size *= s + + # XXX I'm note sure how the memory for self.p[0].buf is handled. I'm + # worried that it is freed when self.p[0] goes out of scope, which may be + # before efield goes out of scope. The fix is to store a reference of + # self.p[0] in efield. This is done below, but doesn't seem to be + # nessisary. -KM + cdef np.ndarray[np.uint8_t,ndim=3,mode='c'] efield + cdef uint8_t *buf + buf = self.p[0].get_buf() + #buf = malloc(size * sizeof(uint8_t)) + efield = np.PyArray_SimpleNewFromData(nd, shape_c, np.NPY_UINT8, buf) + efield.flags.writeable = False + + # Store reference to self to keep it in scope and from freeing + # memory. + # This is the wrong place for class definition if this ends up being + # needed (should be top level). + #class pyobj_array(np.ndarray): + # pass + #efield = efield.view(pyobj_array) + #pyobj_array._ref_to_memory = self + + free(shape_c) + return (t0, nt, efield) + + ############################################## Assembler ######################################### diff --git a/ch_vdif_assembler_pxd.pxd b/ch_vdif_assembler_pxd.pxd index 028b647..0a094ca 100644 --- a/ch_vdif_assembler_pxd.pxd +++ b/ch_vdif_assembler_pxd.pxd @@ -1,4 +1,4 @@ -from libc.stdint cimport int32_t, int64_t +from libc.stdint cimport int32_t, int64_t, uint8_t from libcpp cimport bool from libcpp.string cimport string @@ -24,6 +24,7 @@ cdef extern from "ch_vdif_assembler_cython.hpp" namespace "ch_vdif_assembler": int nt void fill_efield(float complex *efield, int32_t *mask) except + + uint8_t* get_buf() except + cdef cppclass cython_assembler: