Merge pull request #218 from ledatelescope/dcp-sigproc-detect-update

Updates to detect.py, sigproc.py and quantize.py blocks

.mypy.ini

@@ -0,0 +1,4 @@
+[mypy]
+
+[mypy-scipy.fftpack.*]
+ignore_missing_imports = True

Makefile.in

@@ -18,6 +18,19 @@ libbifrost:
 	$(MAKE) -C $(SRC_DIR) all
 .PHONY: libbifrost
 
+check:
+ifeq ($(HAVE_PYTHON),1)
+	MYPYPATH=$(BIFROST_PYTHON_DIR) mypy --follow-imports=silent \
+	    python/bifrost/blocks/detect.py \
+	    python/bifrost/blocks/quantize.py \
+	    python/bifrost/blocks/unpack.py \
+	    python/bifrost/sigproc.py \
+	    python/bifrost/sigproc2.py \
+	    test/test_sigproc.py \
+	    testbench/test_fft_detect.py
+endif
+.PHONY: check
+
 test:
 	#$(MAKE) -C $(SRC_DIR) test
 ifeq ($(HAVE_PYTHON),1)

python/bifrost/blocks/detect.py

@@ -30,12 +30,13 @@
 from bifrost.DataType import DataType
 
 from copy import deepcopy
+from typing import Optional, Union
 
 from bifrost import telemetry
 telemetry.track_module()
 
 class DetectBlock(TransformBlock):
-    def __init__(self, iring, mode, axis=None,
+    def __init__(self, iring, mode: str, axis: Optional[Union[int,str]] = None,
                  *args, **kwargs):
         super(DetectBlock, self).__init__(iring, *args, **kwargs)
         self.specified_axis = axis
@@ -56,7 +57,7 @@ def on_sequence(self, iseq):
               self.mode != 'scalar' and
               'pol' in itensor['labels']):
             self.axis = itensor['labels'].index('pol')
-        elif isinstance(self.axis, str):
+        elif isinstance(self.axis, basestring):
             self.axis = itensor['labels'].index(self.axis)
         # Note: axis may be None here, which indicates single-pol mode
         ohdr = deepcopy(ihdr)
@@ -65,8 +66,10 @@ def on_sequence(self, iseq):
             self.npol = otensor['shape'][self.axis]
             if self.npol not in [1, 2]:
                 raise ValueError("Axis must have length 1 or 2")
-            if self.mode == 'stokes' and self.npol == 2:
+            if (self.mode == 'stokes' or self.mode == 'coherence') and self.npol == 2:
                 otensor['shape'][self.axis] = 4
+            if self.mode == 'stokes_i' and self.npol == 2:
+                otensor['shape'][self.axis] = 1
             if 'labels' in otensor:
                 otensor['labels'][self.axis] = 'pol'
         else:
@@ -109,32 +112,47 @@ def on_data(self, ispan, ospan):
                 b(%s) = -2*xy.imag;
                 """ % (inds_[0], inds_[1],
                        inds_[0], inds_[1], inds_[2], inds_[3])
+            elif self.mode == 'stokes_i':
+                func = """
+                Complex<b_type> x = a(%s);
+                Complex<b_type> y = a(%s);
+                auto xx = x.mag2();
+                auto yy = y.mag2();
+                b(%s) = xx + yy;
+                """ % (inds_[0], inds_[1],
+                       inds_[0])
+            elif self.mode == 'coherence':
+                func = """
+                Complex<b_type> x = a(%s);
+                Complex<b_type> y = a(%s);
+                auto xx = x.mag2();
+                auto yy = y.mag2();
+                auto xy = x.conj()*y;
+                b(%s) = xx;
+                b(%s) = yy;
+                b(%s) = xy.real;
+                b(%s) = xy.imag;
+                """ % (inds_[0], inds_[1],
+                       inds_[0], inds_[1], inds_[2], inds_[3])
             bf_map(func, shape=shape, axis_names=inds,
                    data={'a': ispan.data, 'b': ospan.data})
 
-def detect(iring, mode, axis=None, *args, **kwargs):
+def detect(iring, mode: str, axis: Optional[Union[int,str]] = None, *args, **kwargs):
     """Apply square-law detection to create polarization products.
-
     Args:
         iring (Ring or Block): Input data source.
         mode (string):
-
            ``'scalar': x   -> real x.x*``
-
            ``'jones':  x,y -> complex x.x* + 1j*y.y*, x.y*``
-
            ``'stokes': x,y -> real I, Q, U, V``
-
+           ``'stokes_I': x,y -> x.x* + y.y* (Stokes I)``
         axis: Integer or string specifying the polarization axis. Defaults to
                 'pol'. Not used if mode = 'scalar'.
         *args: Arguments to ``bifrost.pipeline.TransformBlock``.
         **kwargs: Keyword Arguments to ``bifrost.pipeline.TransformBlock``.
-
     **Tensor semantics**::
-
         Input:  [..., 'pol', ...], dtype = any complex, space = CUDA
         Output: [..., 'pol', ...], dtype = real or complex, space = CUDA
-
     Returns:
         DetectBlock: A new block instance.
     """

python/bifrost/blocks/quantize.py

@@ -28,22 +28,26 @@
 from bifrost.quantize import quantize as bf_quantize
 from bifrost.pipeline import TransformBlock
 from bifrost.DataType import DataType
+from bifrost.ring2 import Ring, ReadSequence, ReadSpan, WriteSpan
 
 from copy import deepcopy
+import numpy as np
+
+from typing import Any, Dict, Union, Tuple
 
 from bifrost import telemetry
 telemetry.track_module()
 
 class QuantizeBlock(TransformBlock):
-    def __init__(self, iring, dtype, scale=1.,
+    def __init__(self, iring: Ring, dtype: Union[str,np.dtype], scale: float=1.,
                  *args, **kwargs):
         super(QuantizeBlock, self).__init__(iring, *args, **kwargs)
         self.dtype = dtype
         self.scale = scale
-    def define_valid_input_spaces(self):
+    def define_valid_input_spaces(self) -> Tuple[str]:
         """Return set of valid spaces (or 'any') for each input"""
-        return ('system',)
-    def on_sequence(self, iseq):
+        return ('any',)
+    def on_sequence(self, iseq: ReadSequence) -> Dict[str,Any]:
         ihdr = iseq.header
         ohdr = deepcopy(ihdr)
         itype = DataType(ihdr['_tensor']['dtype'])
@@ -56,12 +60,12 @@ def on_sequence(self, iseq):
             otype = self.dtype
         ohdr['_tensor']['dtype'] = otype
         return ohdr
-    def on_data(self, ispan, ospan):
+    def on_data(self, ispan: ReadSpan, ospan: WriteSpan) -> None:
         idata = ispan.data
         odata = ospan.data
         bf_quantize(idata, odata, self.scale)
 
-def quantize(iring, dtype, scale=1., *args, **kwargs):
+def quantize(iring: Ring, dtype: Union[str,np.dtype], scale: float=1., *args, **kwargs) -> QuantizeBlock:
     """Apply a requantization of bit depth for the data.
 
     Args:
@@ -73,10 +77,10 @@ def quantize(iring, dtype, scale=1., *args, **kwargs):
 
     **Tensor semantics**::
 
-        Input:  [...], dtype = [c]f32, space = SYSTEM
-        Output: [...], dtype = any (complex) integer type, space = SYSTEM
+        Input:  [...], dtype = [c]f32, space = SYSTEM or CUDA
+        Output: [...], dtype = any (complex) integer type, space = SYSTEM or CUDA
 
     Returns:
         QuantizeBlock: A new block instance.
     """
-    return QuantizeBlock(iring, dtype, *args, **kwargs)
+    return QuantizeBlock(iring, dtype, scale, *args, **kwargs)

python/bifrost/blocks/unpack.py

@@ -33,7 +33,7 @@
 from copy import deepcopy
 import numpy as np
 
-from typing import Any, Dict, Tuple, Union
+from typing import Any, Dict, Union, Tuple
 
 from bifrost import telemetry
 telemetry.track_module()
@@ -46,7 +46,7 @@ def __init__(self, iring: Ring, dtype: Union[str,np.dtype], align_msb: bool=Fals
         self.align_msb = align_msb
     def define_valid_input_spaces(self) -> Tuple[str]:
         """Return set of valid spaces (or 'any') for each input"""
-        return ('system',)
+        return ('any',)
     def on_sequence(self, iseq: ReadSequence) -> Dict[str,Any]:
         ihdr = iseq.header
         ohdr = deepcopy(ihdr)
@@ -76,8 +76,8 @@ def unpack(iring: Ring, dtype: Union[str,np.dtype], *args, **kwargs) -> UnpackBl
 
     **Tensor semantics**::
 
-        Input:  [...], dtype = one of: i/u2, i/u4, ci2, ci4, space = SYSTEM
-        Output: [...], dtype = i8 or ci8 (matching input), space = SYSTEM
+        Input:  [...], dtype = one of: i/u2, i/u4, ci2, ci4, space = SYSTEM or CUDA
+        Output: [...], dtype = i8 or ci8 (matching input), space = SYSTEM or CUDA
 
     Returns:
         UnpackBlock: A new block instance.

python/bifrost/sigproc.py

@@ -114,8 +114,17 @@
                            6:  'GBT',
                            7:  'GMRT',
                            8:  'Effelsberg',
+                           9: 'Effelsberg LOFAR',
+                           11: 'Unknown',
+                           12: 'MWA',
+                           20: 'CHIME',
                            52: 'LWA-OV',
-                           53: 'LWA-SV'})
+                           53: 'LWA-SV',
+                           64: 'MeerKAT',
+                           65: 'KAT-7',
+                           82: 'eMerlin'})
+
+
 #the machine_id parameter names' translation
 _MACHINES = defaultdict(lambda: 'unknown',
                         {0:  'FAKE',
@@ -127,6 +136,8 @@
                          6:  'SCAMP',
                          7:  'GMRTFB',
                          8:  'PULSAR2000',
+                         9:  'UNKNOWN',
+                         20: 'CHIME',
                          52: 'LWA-DP',
                          53: 'LWA-ADP'})
 
@@ -248,7 +259,7 @@ def pack(data: np.ndarray, nbit: int) -> np.ndarray:
         outdata += data[index::8 // nbit] // (2**nbit)**index
     return outdata
 
-def _write_data(data: np.ndarray, nbit: int, file_object: IO[bytes]) -> np.ndarray:
+def _write_data(data: np.ndarray, nbit: int, file_object: IO[bytes]):
     """Writes given data to an open file, also packing if needed"""
     file_object.seek(0, 2)
     if nbit < 8:

python/bifrost/sigproc2.py

@@ -113,11 +113,17 @@
                            6:  'GBT',
                            7:  'GMRT',
                            8:  'Effelsberg',
-                           9:  'ATA',
+                           9:  'Effelsberg LOFAR',
+                           11: 'Unknown',
+                           12: 'MWA',
+                           20: 'CHIME',
                            10: 'UTR-2',
                            11: 'LOFAR',
                            52: 'LWA-OV',
-                           53: 'LWA-SV'})
+                           53: 'LWA-SV',
+                           64: 'MeerKAT',
+                           65: 'KAT-7',
+                           82: 'eMerlin'})
 _machines   = defaultdict(lambda: 'unknown',
                           {0:  'FAKE',
                            1:  'PSPM',
@@ -128,6 +134,8 @@
                            6:  'SCAMP',
                            7:  'GMRTFB', # aka GBT Pulsar Spigot, SPIGOT
                            8:  'PULSAR2000',
+                           9:  'UNKNOWN',
+                           20: 'CHIME',
                            11: 'BG/P',
                            12: "PDEV",
                            20: 'GUPPI',
@@ -275,10 +283,11 @@ def open(self, filename: str) -> "SigprocFile":
         self.signed = bool(self.header['signed'])
         if self.nbit >= 8:
             if self.signed:
-                self.dtype  = { 8: np.int8,
-                               16: np.int16,
-                               32: np.float32,
-                               64: np.float64}[self.nbit]
+                self.dtype : Optional[type] = {
+                    8: np.int8,
+                    16: np.int16,
+                    32: np.float32,
+                    64: np.float64}[self.nbit]
             else:
                 self.dtype  = { 8: np.uint8,
                                16: np.uint16,