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