def tensor(self): # TODO: This shouldn't be public
if self._tensor is not None:
return self._tensor
header = self.header
shape = header['_tensor']['shape']
ringlet_shape, frame_shape = split_shape(shape)
nringlet = reduce(lambda x, y: x * y, ringlet_shape, 1)
frame_nelement = reduce(lambda x, y: x * y, frame_shape, 1)
dtype = header['_tensor']['dtype']
try:
dtype = dtype.decode()
except AttributeError:
# Python2 catch
pass
nbit = DataType(dtype).itemsize_bits
assert(nbit % 8 == 0)
frame_nbyte = frame_nelement * nbit // 8
self._tensor = {}
self._tensor['dtype'] = DataType(dtype)
self._tensor['ringlet_shape'] = ringlet_shape
self._tensor['nringlet'] = nringlet
self._tensor['frame_shape'] = frame_shape
self._tensor['frame_nbyte'] = frame_nbyte
self._tensor['dtype_nbyte'] = nbit // 8
return self._tensor
def header_transform(hdr, new_dtype=dtype):
tensor = hdr['_tensor']
old_dtype = tensor['dtype']
old_itemsize = DataType(old_dtype).itemsize
new_itemsize = DataType(new_dtype).itemsize
old_axissize = old_itemsize * tensor['shape'][-1]
if old_axissize % new_itemsize:
raise ValueError("New type not compatible with data shape")
tensor['shape'][-1] = old_axissize // new_itemsize
tensor['dtype'] = dtype
return hdr
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 data_view(self, dtype=np.uint8, shape=-1):
itemsize = DataType(dtype).itemsize
assert (self.size % itemsize == 0)
assert (self.stride % itemsize == 0)
data_ptr = self._data_ptr
#if self.sequence.ring.space == 'cuda':
# # TODO: See if can wrap this in something like PyCUDA's GPUArray
# # Ideally actual GPUArray, but it doesn't appear to support wrapping pointers
# # Could also try writing a custom GPUArray implem for this purpose
# return data_ptr
span_size = self.size
nringlet = self.nringlet
#print("******", span_size, stride, nringlet)
#BufferType = c_byte*(span_size*self.stride)
# TODO: We should really map the actual ring memory space and index
# it with offset rather than mapping from the current pointer.
_shape = (nringlet, span_size // itemsize)
strides = (self.stride, itemsize) if nringlet > 1 else None
space = self.ring.space
data_array = ndarray(shape=_shape,
strides=strides,
buffer=data_ptr,
dtype=dtype,
space=space)
# Note: This is a non-standard attribute
#data_array.flags['SPACE'] = space
if not self.writeable:
data_array.flags['WRITEABLE'] = False
if shape != -1:
# TODO: Check that this still wraps the same memory
data_array = data_array.reshape(shape)
return data_array
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
def on_sequence(self, iseq):
ihdr = iseq.header
ohdr = deepcopy(ihdr)
itype = DataType(ihdr['_tensor']['dtype'])
self.itype = itype
# Allow user to pass nbit instead of explicit dtype
if isinstance(self.dtype, int):
nbit = self.dtype
otype = itype.as_nbit(nbit)
else:
otype = self.dtype
ohdr['_tensor']['dtype'] = otype
return ohdr
def view(self, dtype=None, type_=None):
if type_ is not None:
dtype = type
type_type = type(int) # HACK to form an instance of 'type' (very confusing)
if isinstance(dtype, type_type) and issubclass(dtype, np.ndarray):
return super(ndarray, self).view(dtype)
else:
# TODO: Endianness changes are not supported here
# Consider building byteorder into DataType
dtype_bf = DataType(dtype)
dtype_np = np.dtype(dtype_bf.as_numpy_dtype())
v = super(ndarray, self).view(dtype_np)
v.bf.dtype = dtype_bf
v._update_BFarray()
return v
def __array_finalize__(self, obj):
if obj is None:
# Already initialized self.bf in __new__
return
# Initialize as view of existing array
if isinstance(obj, ndarray):
# Copy metadata from existing bf.ndarray
self.bf = BFArrayInfo(obj.bf.space, obj.bf.dtype,
obj.bf.native, obj.bf.conjugated)
else:
# Generate metadata from existing np.ndarray
#*space = str(Space(raw_get_space(obj.ctypes.data)))
# Note: Assumes that any existing np.ndarray is in system space
space = 'system'
#dtype = str(DataType(obj.dtype))
# **TODO: Decide on bf.dtype being DataType vs. string (and same for space)
dtype = DataType(obj.dtype)
native = obj.dtype.isnative
conjugated = False
self.bf = BFArrayInfo(space, dtype, native, conjugated)
self._update_BFarray()
def on_sequence(self, iseq):
ihdr = iseq.header
itensor = ihdr['_tensor']
axnames = tuple(itensor['labels'])
shape = itensor['shape']
scales = itensor['scales']
units = itensor['units']
ndim = len(shape)
dtype = DataType(itensor['dtype'])
nchan = shape[-1]
sample_time = convert_units(scales[-2][1], units[-2], 's')
sample_rate = int(round(1. / sample_time))
frame_nbyte = nchan * dtype.itemsize
ohdr = {
'audio_fmt': 1, # 1 => PCM (linear quantization, uncompressed)
'nchan': nchan,
'sample_rate': sample_rate,
'byte_rate': sample_rate * frame_nbyte,
'block_align': frame_nbyte,
'nbit': dtype.itemsize_bits
}
filename = os.path.join(self.path, ihdr['name'])
if ndim == 2 and axnames[-2] == 'time':
self.ofile = open(filename + '.wav', 'wb')
wav_write_header(self.ofile, ohdr)
elif ndim == 3 and axnames[-2] == 'time':
nfile = shape[-3]
filenames = [filename + '.%09i.tim' % i for i in range(nfile)]
self.ofiles = [open(fname + '.wav', 'wb') for fname in filenames]
for ofile in self.ofiles:
wav_write_header(ofile, ohdr)
else:
raise ValueError("Incompatible axes: " + str(axnames))
def on_data(self, ispan, ospan):
idata = ispan.data
odata = ospan.data
itype = DataType(idata.dtype)
otype = DataType(odata.dtype)
if self.ifmt == 'matrix' and self.ofmt == 'matrix':
# Make a full-matrix copy of the lower-only input matrix
# odata[t,c,i,p,j,q] = idata[t,c,i,p,j,q] (lower filled only)
shape_nopols = list(idata.shape)
del shape_nopols[5]
del shape_nopols[3]
idata = idata.view(itype.as_vector(2))
odata = odata.view(otype.as_vector(2))
bf.map('''
bool in_lower_triangle = (i > j);
if( in_lower_triangle ) {
odata(t,c,i,0,j,0) = idata(t,c,i,0,j,0);
odata(t,c,i,1,j,0) = idata(t,c,i,1,j,0);
} else {
auto x = idata(t,c,j,0,i,0);
auto y = idata(t,c,j,1,i,0);
auto x1 = x[1];
x[0] = x[0].conj();
x[1] = y[0].conj();
if( i != j ) {
y[0] = x1.conj();
}
y[1] = y[1].conj();
odata(t,c,i,0,j,0) = x;
odata(t,c,i,1,j,0) = y;
}
''',
shape=shape_nopols,
axis_names=['t', 'c', 'i', 'j'],
data={
'idata': idata,
'odata': odata
})
elif self.ifmt == 'matrix' and self.ofmt == 'storage':
assert (idata.shape[2] <= 2048)
idata = idata.view(itype.as_vector(2))
odata = odata.view(otype.as_vector(4))
# TODO: Support L/R as well as X/Y pols
bf.map('''
// TODO: This only works up to 2048 in single-precision
#define project_triangular(i, j) ((i)*((i)+1)/2 + (j))
int i = int((sqrt(8.f*(b)+1)-1)/2);
int j = b - project_triangular(i, 0);
auto x = idata(t,c,i,0,j,0);
auto y = idata(t,c,i,1,j,0);
if( i == j ) {
x[1] = y[0].conj();
}
idata_type::value_type eye(0, 1);
auto I = (x[0] + y[1]);
auto Q = (x[0] - y[1]);
auto U = (x[1] + y[0]);
auto V = (x[1] - y[0]) * eye;
odata(t,b,c,0) = odata_type(I,Q,U,V);
''',
shape=odata.shape[:-1],
axis_names=['t', 'b', 'c'],
data={
'idata': idata,
'odata': odata
},
block_shape=[64, 8]) # TODO: Tune this
#elif self.ifmt == 'matrix' and self.ofmt == 'triangular':
elif self.ifmt == 'storage' and self.ofmt == 'matrix':
oshape_nopols = list(odata.shape)
del oshape_nopols[5]
del oshape_nopols[3]
idata = idata.view(itype.as_vector(4))
odata = odata.view(otype.as_vector(2))
bf.map('''
bool in_upper_triangle = (i < j);
auto b = in_upper_triangle ? j*(j+1)/2 + i : i*(i+1)/2 + j;
auto IQUV = idata(t,b,c,0);
auto I = IQUV[0], Q = IQUV[1], U = IQUV[2], V = IQUV[3];
idata_type::value_type eye(0, 1);
auto xx = 0.5f*(I + Q);
auto xy = 0.5f*(U - V*eye);
auto yx = 0.5f*(U + V*eye);
auto yy = 0.5f*(I - Q);
if( i == j ) {
xy = yx.conj();
}
if( in_upper_triangle ) {
auto tmp_xy = xy;
xx = xx.conj();
xy = yx.conj();
yx = tmp_xy.conj();
yy = yy.conj();
}
odata(t,c,i,0,j,0) = odata_type(xx, xy);
odata(t,c,i,1,j,0) = odata_type(yx, yy);
''',
shape=oshape_nopols,
axis_names=['t', 'c', 'i', 'j'],
data={
'idata': idata,
'odata': odata
},
block_shape=[64, 8]) # TODO: Tune this
else:
raise NotImplementedError
def on_sequence(self, iseq):
ihdr = iseq.header
itensor = ihdr['_tensor']
axnames = list(itensor['labels'])
shape = list(itensor['shape'])
scales = list(itensor['scales'])
units = list(itensor['units'])
ndim = len(shape)
dtype = DataType(itensor['dtype'])
sigproc_hdr = {}
_copy_item_if_exists(sigproc_hdr, ihdr, 'source_name')
_copy_item_if_exists(sigproc_hdr, ihdr, 'rawdatafile')
_copy_item_if_exists(sigproc_hdr, ihdr, 'az_start')
_copy_item_if_exists(sigproc_hdr, ihdr, 'za_start')
_copy_item_if_exists(sigproc_hdr, ihdr, 'raj', 'src_raj')
_copy_item_if_exists(sigproc_hdr, ihdr, 'dej', 'src_dej')
if 'telescope' in ihdr:
sigproc_hdr['telescope_id'] = sigproc.telescope2id(
ihdr['telescope'])
if 'machine' in ihdr:
sigproc_hdr['machine_id'] = sigproc.machine2id(ihdr['machine'])
_copy_item_if_exists(sigproc_hdr, ihdr, 'telescope_id')
_copy_item_if_exists(sigproc_hdr, ihdr, 'machine_id')
_copy_item_if_exists(sigproc_hdr, ihdr, 'ibeam')
_copy_item_if_exists(sigproc_hdr, ihdr, 'nbeams')
sigproc_hdr['nbits'] = dtype.itemsize_bits
_copy_item_if_exists(sigproc_hdr, ihdr, 'barycentric')
_copy_item_if_exists(sigproc_hdr, ihdr, 'pulsarcentric')
if dtype.is_integer and dtype.is_signed:
sigproc_hdr['signed'] = True
if 'coord_frame' in ihdr:
coord_frame = ihdr['coord_frame']
else:
coord_frame = None
sigproc_hdr['pulsarcentric'] = (coord_frame == 'pulsarcentric')
sigproc_hdr['barycentric'] = (coord_frame == 'barycentric')
filename = os.path.join(self.path, ihdr['name'])
if ndim >= 3 and axnames[-3:] == ['time', 'pol', 'freq']:
self.data_format = 'filterbank'
assert (dtype.is_real)
sigproc_hdr['data_type'] = 1
sigproc_hdr['nifs'] = shape[-2]
sigproc_hdr['nchans'] = shape[-1]
sigproc_hdr['tstart'] = _unix2mjd(scales[-3][0])
sigproc_hdr['tsamp'] = convert_units(scales[-3][1], units[-3], 's')
sigproc_hdr['fch1'] = convert_units(scales[-1][0], units[-1],
'MHz')
sigproc_hdr['foff'] = convert_units(scales[-1][1], units[-1],
'MHz')
if 'refdm' in ihdr:
sigproc_hdr['refdm'] = convert_units(ihdr['refdm'],
ihdr['refdm_units'],
'pc cm^-3')
if ndim == 3:
filename += '.fil'
self.ofile = open(filename, 'wb')
sigproc.write_header(sigproc_hdr, self.ofile)
elif ndim == 4:
if axnames[-4] != 'beam':
raise ValueError("Expected first axis to be 'beam'"
" got '%s'" % axnames[-4])
nbeam = shape[-4]
sigproc_hdr['nbeams'] = nbeam
filenames = [
filename + '.%06iof.%06i.fil' % (b + 1, nbeam)
for b in range(nbeam)
]
self.ofiles = [open(fname, 'wb') for fname in filenames]
for b in range(nbeam):
sigproc_hdr['ibeam'] = b
sigproc.write_header(sigproc_hdr, self.ofiles[b])
else:
raise ValueError("Too many dimensions")
elif ndim >= 2 and 'time' in axnames and 'pol' in axnames:
pol_axis = axnames.index('pol')
if pol_axis != ndim - 1:
# Need to move pol axis
# Note: We support this because it tends to be convenient
# for rest of the pipeline to operate with pol being
# the first dim, and doing the transpose on the fly
# inside this block is unlikely to cost much relative
# to disk perf (and it's free if npol==1).
axnames.append(axnames[pol_axis])
del axnames[pol_axis]
shape.append(shape[pol_axis])
del shape[pol_axis]
scales.append(scales[pol_axis])
del scales[pol_axis]
units.append(units[pol_axis])
del units[pol_axis]
self.pol_axis = pol_axis
self.data_format = 'timeseries'
assert (dtype.is_real)
sigproc_hdr['data_type'] = 2
sigproc_hdr['nchans'] = 1
sigproc_hdr['nifs'] = shape[-2]
sigproc_hdr['tstart'] = _unix2mjd(scales[-2][0])
sigproc_hdr['tsamp'] = convert_units(scales[-2][1], units[-2], 's')
if 'cfreq' in ihdr and 'bw' in ihdr:
sigproc_hdr['fch1'] = convert_units(ihdr['cfreq'],
ihdr['cfreq_units'], 'MHz')
sigproc_hdr['foff'] = convert_units(ihdr['bw'],
ihdr['bw_units'], 'MHz')
# TODO: Write ndim separate output files, each with its own refdm
if ndim == 2:
if 'refdm' in ihdr:
sigproc_hdr['refdm'] = convert_units(
ihdr['refdm'], ihdr['refdm_units'], 'pc cm^-3')
filename += '.tim'
self.ofile = open(filename, 'wb')
sigproc.write_header(sigproc_hdr, self.ofile)
elif ndim == 3:
if axnames[-3] != 'dispersion':
raise ValueError("Expected first axis to be 'dispersion'"
" got '%s'" % axnames[-3])
ndm = shape[-3]
dm0 = scales[-3][0]
ddm = scales[-3][1]
dms = [dm0 + ddm * d for d in range(ndm)]
dms = [convert_units(dm, units[-3], 'pc cm^-3') for dm in dms]
filenames = [filename + '.%09.2f.tim' % dm for dm in dms]
self.ofiles = [open(fname, 'wb') for fname in filenames]
for d, dm in enumerate(dms):
sigproc_hdr['refdm'] = dm
sigproc.write_header(sigproc_hdr, self.ofiles[d])
else:
raise ValueError("Too many dimensions")
elif ndim == 4 and axnames[-3:] == ['pol', 'freq', 'phase']:
self.data_format = 'pulseprofile'
assert (dtype.is_real)
sigproc_hdr['data_type'] = 2
sigproc_hdr['nifs'] = shape[-3]
sigproc_hdr['nchans'] = shape[-2]
sigproc_hdr['nbins'] = shape[-1]
sigproc_hdr['tstart'] = _unix2mjd(scales[-4][0])
sigproc_hdr['tsamp'] = convert_units(scales[-4][1], units[-4], 's')
sigproc_hdr['fch1'] = convert_units(scales[-2][0], units[-2],
'MHz')
sigproc_hdr['foff'] = convert_units(scales[-2][1], units[-2],
'MHz')
if 'refdm' in ihdr:
sigproc_hdr['refdm'] = convert_units(ihdr['refdm'],
ihdr['refdm_units'],
'pc cm^-3')
_copy_item_if_exists(sigproc_hdr, ihdr, 'npuls')
self.filename = filename
self.sigproc_hdr = sigproc_hdr
self.t0 = scales[-4][0]
self.dt = scales[-4][1]
else:
raise ValueError(
"Axis labels do not correspond to a known data format: " +
str(axnames) + "\nKnown formats are:" +
"\n [time, pol, freq]\n [beam, time, pol]\n" +
" [time, pol]\n [dispersion, time, pol]\n" +
" [pol, freq, phase]")
def on_sequence(self, iseq):
ihdr = iseq.header
itensor = ihdr['_tensor']
axnames = tuple(itensor['labels'])
shape = itensor['shape']
scales = itensor['scales']
units = itensor['units']
ndim = len(shape)
dtype = DataType(itensor['dtype'])
sigproc_hdr = {}
copy_item_if_exists(sigproc_hdr, ihdr, 'source_name')
copy_item_if_exists(sigproc_hdr, ihdr, 'rawdatafile')
copy_item_if_exists(sigproc_hdr, ihdr, 'az_start')
copy_item_if_exists(sigproc_hdr, ihdr, 'za_start')
copy_item_if_exists(sigproc_hdr, ihdr, 'raj', 'src_raj')
copy_item_if_exists(sigproc_hdr, ihdr, 'dej', 'src_dej')
if 'telescope' in ihdr:
sigproc_hdr['telescope_id'] = sigproc.telescope2id(
ihdr['telescope'])
if 'machine' in ihdr:
sigproc_hdr['machine_id'] = sigproc.machine2id(ihdr['machine'])
copy_item_if_exists(sigproc_hdr, ihdr, 'telescope_id')
copy_item_if_exists(sigproc_hdr, ihdr, 'machine_id')
copy_item_if_exists(sigproc_hdr, ihdr, 'ibeam')
copy_item_if_exists(sigproc_hdr, ihdr, 'nbeams')
sigproc_hdr['nbits'] = dtype.itemsize_bits
copy_item_if_exists(sigproc_hdr, ihdr, 'barycentric')
copy_item_if_exists(sigproc_hdr, ihdr, 'pulsarcentric')
if dtype.is_integer and dtype.is_signed:
sigproc_hdr['signed'] = True
if 'coord_frame' in ihdr:
coord_frame = ihdr['coord_frame']
else:
coord_frame = None
sigproc_hdr['pulsarcentric'] = (coord_frame == 'pulsarcentric')
sigproc_hdr['barycentric'] = (coord_frame == 'barycentric')
filename = os.path.join(self.path, ihdr['name'])
if ndim >= 3 and axnames[-3:] == ('time', 'pol', 'freq'):
self.data_format = 'filterbank'
assert (dtype.is_real)
sigproc_hdr['data_type'] = 1
sigproc_hdr['nifs'] = shape[-2]
sigproc_hdr['nchans'] = shape[-1]
sigproc_hdr['tstart'] = unix2mjd(scales[-3][0])
sigproc_hdr['tsamp'] = convert_units(scales[-3][1], units[-3], 's')
sigproc_hdr['fch1'] = convert_units(scales[-1][0], units[-1],
'MHz')
sigproc_hdr['foff'] = convert_units(scales[-1][1], units[-1],
'MHz')
if 'refdm' in ihdr:
sigproc_hdr['refdm'] = convert_units(ihdr['refdm'],
ihdr['refdm_units'],
'pc cm^-3')
if ndim == 3:
filename += '.fil'
self.ofile = open(filename, 'wb')
sigproc.write_header(sigproc_hdr, self.ofile)
elif ndim == 4:
if axnames[-4] != 'beam':
raise ValueError("Expected first axis to be 'beam'")
nbeam = shape[-4]
sigproc_hdr['nbeams'] = nbeam
filenames = [
filename + '.%06iof.%06i.fil' % (b + 1, nbeam)
for b in xrange(nbeam)
]
self.ofiles = [open(fname, 'wb') for fname in filenames]
for b in xrange(nbeam):
sigproc_hdr['ibeam'] = b
sigproc.write_header(sigproc_hdr, self.ofiles[b])
else:
raise ValueError("Too many dimensions")
elif ndim >= 2 and axnames[-2:] == ('time', 'pol'):
self.data_format = 'timeseries'
assert (dtype.is_real)
sigproc_hdr['data_type'] = 2
sigproc_hdr['nchans'] = 1
sigproc_hdr['nifs'] = shape[-2]
sigproc_hdr['tstart'] = unix2mjd(scales[-2][0])
sigproc_hdr['tsamp'] = convert_units(scales[-2][1], units[-2], 's')
if 'cfreq' in ihdr and 'bw' in ihdr:
sigproc_hdr['fch1'] = convert_units(ihdr['cfreq'],
ihdr['cfreq_units'], 'MHz')
sigproc_hdr['foff'] = convert_units(ihdr['bw'],
ihdr['bw_units'], 'MHz')
# TODO: Write ndim separate output files, each with its own refdm
if ndim == 2:
if 'refdm' in ihdr:
sigproc_hdr['refdm'] = convert_units(
ihdr['refdm'], ihdr['refdm_units'], 'pc cm^-3')
filename += '.tim'
self.ofile = open(filename, 'wb')
sigproc.write_header(sigproc_hdr, self.ofile)
elif ndim == 3:
if axnames[-3] != 'dispersion measure':
raise ValueError(
"Expected first axis to be 'dispersion measure'")
ndm = shape[-3]
dm0 = scales[-3][0]
ddm = scales[-3][1]
dms = [dm0 + ddm * d for d in xrange(ndm)]
dms = [convert_units(dm, units[-3], 'pc cm^-3') for dm in dms]
filenames = [filename + '.%09.2f.tim' % dm for dm in dms]
self.ofiles = [open(fname, 'wb') for fname in filenames]
for d, dm in enumerate(dms):
sigproc_hdr['refdm'] = dm
sigproc.write_header(sigproc_hdr, self.ofiles[d])
else:
raise ValueError("Too many dimensions")
elif ndim == 4 and axnames[-3:] == ('pol', 'freq', 'phase'):
self.data_format = 'pulseprofile'
assert (dtype.is_real)
sigproc_hdr['data_type'] = 2
sigproc_hdr['nifs'] = shape[-3]
sigproc_hdr['nchans'] = shape[-2]
sigproc_hdr['nbins'] = shape[-1]
sigproc_hdr['tstart'] = unix2mjd(scales[-4][0])
sigproc_hdr['tsamp'] = convert_units(scales[-4][1], units[-4], 's')
sigproc_hdr['fch1'] = convert_units(scales[-2][0], units[-2],
'MHz')
sigproc_hdr['foff'] = convert_units(scales[-2][1], units[-2],
'MHz')
if 'refdm' in ihdr:
sigproc_hdr['refdm'] = convert_units(ihdr['refdm'],
ihdr['refdm_units'],
'pc cm^-3')
copy_item_if_exists(sigproc_hdr, ihdr, 'npuls')
self.filename = filename
self.sigproc_hdr = sigproc_hdr
self.t0 = scales[-4][0]
self.dt = scales[-4][1]
else:
raise ValueError(
"Axis labels do not correspond to a known data format: " +
str(axnames))
def __new__(cls, base=None, space=None, shape=None, dtype=None,
buffer=None, offset=0, strides=None,
native=None, conjugated=None):
if isinstance(shape, int):
shape = [shape]
ownbuffer = None
if base is not None:
if (shape is not None or
# dtype is not None or
buffer is not None or
offset != 0 or
strides is not None or
native is not None):
raise ValueError('Invalid combination of arguments when base '
'is specified')
if 'cupy' in sys.modules:
from cupy import ndarray as cupy_ndarray
if isinstance(base, cupy_ndarray):
return ndarray.__new__(cls,
space='cuda',
buffer=int(base.data),
shape=base.shape,
dtype=base.dtype,
strides=base.strides,
native=np.dtype(base.dtype).isnative)
if 'pycuda' in sys.modules:
from pycuda.gpuarray import GPUArray as pycuda_GPUArray
if isinstance(base, pycuda_GPUArray):
return ndarray.__new__(cls,
space='cuda',
buffer=int(base.gpudata),
shape=base.shape,
dtype=base.dtype,
strides=base.strides,
native=np.dtype(base.dtype).isnative)
if dtype is not None:
dtype = DataType(dtype)
if space is None and dtype is None:
if not isinstance(base, np.ndarray):
base = np.asarray(base)
# TODO: This may not be a good idea
# Create view of base array
obj = base.view(cls) # Note: This calls obj.__array_finalize__
# Allow conjugated to be redefined
if conjugated is not None:
obj.bf.conjugated = conjugated
obj._update_BFarray()
else:
if not isinstance(base, np.ndarray):
# Convert base to np.ndarray
if dtype is not None:
base = np.array(base,
dtype=DataType(dtype).as_numpy_dtype())
else:
base = np.array(base)
if not isinstance(base, ndarray) and dtype is not None:
base = base.astype(dtype.as_numpy_dtype())
base = ndarray(base) # View base as bf.ndarray
if dtype is not None and base.bf.dtype != dtype:
raise TypeError('Unable to convert type %s to %s during '
'array construction' %
(base.bf.dtype, dtype))
#base = base.view(cls
#if dtype is not None:
# base = base.astype(DataType(dtype).as_numpy_dtype())
if conjugated is None:
conjugated = base.bf.conjugated
# Create copy of base array
obj = ndarray.__new__(cls,
space=space,
shape=base.shape,
dtype=base.bf.dtype,
strides=base.strides,
native=base.bf.native,
conjugated=conjugated)
copy_array(obj, base)
else:
# Create new array
if dtype is None:
dtype = 'f32' # Default dtype
dtype = DataType(dtype)
if native is None:
native = True # Default byteorder
if conjugated is None:
conjugated = False # Default unconjugated
if strides is None:
#itemsize = dtype.itemsize
itemsize_bits = dtype.itemsize_bits
# HACK to support 'packed' arrays, by folding the last
# dimension of the shape into the dtype.
# TODO: Consider using bit strides when dtype < 8 bits
# It's hacky, but it may be worth it
if itemsize_bits < 8:
pack_factor = 8 // itemsize_bits
if shape[-1] % pack_factor != 0 or not len(shape):
raise ValueError("Array cannot be packed")
shape = list(shape)
shape[-1] //= pack_factor
itemsize = 1
else:
itemsize = itemsize_bits // 8
if len(shape):
# This magic came from http://stackoverflow.com/a/32874295
strides = (itemsize *
np.r_[1, np.cumprod(shape[::-1][:-1],
dtype=np.int64)][::-1])
strides = tuple(strides)
else:
strides = tuple()
nbyte = strides[0] * shape[0] if len(shape) else itemsize
if buffer is None:
# Allocate new buffer
if space is None:
space = 'system' # Default space
if shape is None:
raise ValueError('Either buffer or shape must be '
'specified')
ownbuffer = raw_malloc(nbyte, space)
buffer = ownbuffer
else:
if space is None:
#space = _get(_bf.bfGetSpace(buffer))
# TODO: raw_get_space should probably return string, and needs a better name
space = str(Space(raw_get_space(buffer)))
# TODO: Should move np.dtype() into as_numpy_dtype?
dtype_np = np.dtype(dtype.as_numpy_dtype())
if not native:
dtype_np = dtype_np.newbyteorder()
data_buffer = _address_as_buffer(buffer, nbyte)
obj = np.ndarray.__new__(cls, shape, dtype_np,
data_buffer, offset, strides)
obj.bf = BFArrayInfo(space, dtype, native, conjugated, ownbuffer)
obj._update_BFarray()
return obj