def on_sequence(self, iseq): ihdr = iseq.header itensor = ihdr['_tensor'] itype = DataType(itensor['dtype']) if not itype.is_complex: raise TypeError("Input data must be complex") self.axis = self.specified_axis if 'labels' not in itensor.keys() and self.axis is None: raise TypeError( "Polarization (pol) index must be labelled, or axis must be set manually" ) elif (self.axis is None and self.mode != 'scalar' and 'pol' in itensor['labels']): self.axis = itensor['labels'].index('pol') elif isinstance(self.axis, str): self.axis = itensor['labels'].index(self.axis) # Note: axis may be None here, which indicates single-pol mode ohdr = deepcopy(ihdr) otensor = ohdr['_tensor'] if self.axis is not None: 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: otensor['shape'][self.axis] = 4 if 'labels' in otensor: otensor['labels'][self.axis] = 'pol' else: self.npol = 1 if self.mode == 'jones' and self.npol == 2: otype = itype else: otype = itype.as_real() otensor['dtype'] = otype.as_floating_point() return ohdr
def on_sequence(self, iseq): ihdr = iseq.header itensor = ihdr['_tensor'] # TODO: DataType cast should be done inside ring2 # **This tensor stuff generally needs to be cleaned up itype = DataType(itensor['dtype']) # TODO: This is slightly hacky; it needs to emulate the type casting # that Bifrost does internally for the FFT. itype = itype.as_floating_point() # Get axis indices, allowing for lookup-by-label self.axes = [ itensor['labels'].index(axis) if isinstance(axis, basestring) else axis for axis in self.specified_axes ] axes = self.axes shape = [itensor['shape'][ax] for ax in axes] otype = itype.as_real() if self.real_output else itype.as_complex() ohdr = deepcopy(ihdr) otensor = ohdr['_tensor'] otensor['dtype'] = str(otype) if itype.is_real and otype.is_complex: self.mode = 'r2c' elif itype.is_complex and otype.is_real: self.mode = 'c2r' else: self.mode = 'c2c' frame_axis = itensor['shape'].index(-1) if frame_axis in axes: raise KeyError( "Cannot transform frame axis; reshape the data stream first") # Adjust output shape for real transforms if self.mode == 'r2c': otensor['shape'][axes[-1]] //= 2 otensor['shape'][axes[-1]] += 1 elif self.mode == 'c2r': otensor['shape'][axes[-1]] -= 1 otensor['shape'][axes[-1]] *= 2 shape[-1] -= 1 shape[-1] *= 2 for i, (ax, length) in enumerate(zip(axes, shape)): if 'units' in otensor: units = otensor['units'][ax] otensor['units'][ax] = transform_units(units, -1) if 'scales' in otensor: otensor['scales'][ax][0] = 0 # TODO: Is this OK? scale = otensor['scales'][ax][1] otensor['scales'][ax][1] = 1. / (scale * length) if 'labels' in otensor and self.axis_labels is not None: otensor['labels'][ax] = self.axis_labels[i] self.nframe = 0 self.istride = 0 self.ostride = 0 return ohdr