def reduce(idata, odata, op='sum'):
if op not in REDUCE_MAP:
raise ValueError("Invalid reduce op: " + str(op))
op = REDUCE_MAP[op]
_check(_bf.bfReduce(asarray(idata).as_BFarray(),
asarray(odata).as_BFarray(),
op))
def execute(self, idata, odata):
# TODO: Work out how to integrate CUDA stream
_check(
_bf.bfFirExecute(self.obj,
asarray(idata).as_BFarray(),
asarray(odata).as_BFarray()))
return odata
def matmul(self, alpha, a, b, beta, c):
"""Computes:
c = alpha*a.b + beta*c
or if b is None:
c = alpha*a.a^H + beta*c
or if a is None:
c = alpha*b^H.b + beta*c
where '.' is matrix product and '^H' is Hermitian transpose.
Multi-dimensional semantics are the same as numpy.matmul:
The last two dims represent the matrix, and all other dims are
used as batch dims to be matched or broadcast between a and b.
"""
if alpha is None:
alpha = 1.
if beta is None:
beta = 0.
beta = float(beta)
alpha = float(alpha)
a_array = asarray(a).as_BFarray() if a is not None else None
b_array = asarray(b).as_BFarray() if b is not None else None
c_array = asarray(c).as_BFarray()
_check(_bf.bfLinAlgMatMul(self.obj,
alpha,
a_array,
b_array,
beta,
c_array))
return c
def transpose(dst, src, axes=None):
if axes is None:
axes = reversed(range(len(dst.shape)))
dst_bf = asarray(dst).as_BFarray()
src_bf = asarray(src).as_BFarray()
array_type = ctypes.c_int * src.ndim
axes_array = array_type(*axes)
_check(_bf.bfTranspose(src_bf, dst_bf, axes_array))
def execute_workspace(self, iarray, oarray, workspace_ptr, workspace_size,
inverse=False):
_check(_bf.bfFftExecute(
self.obj,
asarray(iarray).as_BFarray(),
asarray(oarray).as_BFarray(),
inverse,
workspace_ptr, workspace_size))
return oarray
def execute(self, idata, odata, negative_delays=False):
# TODO: Work out how to integrate CUDA stream
psize = None
_check(
_bf.bfFdmtExecute(self.obj,
asarray(idata).as_BFarray(),
asarray(odata).as_BFarray(), negative_delays,
None, psize))
return odata
def init(self, iarray, oarray, axes=None, apply_fftshift=False):
if isinstance(axes, int):
axes = [axes]
ndim = len(axes)
if axes is not None:
axes_type = ctypes.c_int * ndim
axes = axes_type(*axes)
self.workspace_size = _get(_bf.bfFftInit, self.obj,
asarray(iarray).as_BFarray(),
asarray(oarray).as_BFarray(), ndim, axes,
apply_fftshift)
def execute_workspace(self,
idata,
odata,
workspace_ptr,
workspace_size,
negative_delays=False):
size = _bf.BFsize(workspace_size)
_check(
_bf.bfFdmtExecute(self.obj,
asarray(idata).as_BFarray(),
asarray(odata).as_BFarray(), negative_delays,
workspace_ptr, size))
return odata
def init(self, coeffs, decim=1, space='cuda'):
space = _string2space(space)
psize = None
_check(
_bf.bfFirInit(self.obj,
asarray(coeffs).as_BFarray(), decim, space, 0,
psize))
def _convert_to_array(arg):
if _is_literal(arg):
arr = np.array(arg)
if isinstance(arg, (int, long)) and -(1 << 31) <= arg < (1 << 31):
arr = arr.astype(np.int32)
# TODO: Any way to decide when these should be double-precision?
elif isinstance(arg, float):
arr = arr.astype(np.float32)
elif isinstance(arg, complex):
arr = arr.astype(np.complex64)
arr.flags['WRITEABLE'] = False
arg = arr
return asarray(arg)
def set_coeffs(self, coeffs):
_check(_bf.bfFirSetCoeffs(self.obj, asarray(coeffs).as_BFarray()))
def unpack(src, dst, align_msb=False):
src_bf = asarray(src).as_BFarray()
dst_bf = asarray(dst).as_BFarray()
_check(_bf.bfUnpack(src_bf, dst_bf, align_msb))
return dst
def set_kernels(self, kernels):
_check(_bf.bfRomeinSetKernels(self.obj, asarray(kernels).as_BFarray()))
def get_workspace_size(self, idata, odata):
return _get(_bf.bfFdmtExecute, self.obj,
asarray(idata).as_BFarray(),
asarray(odata).as_BFarray(), False, None)
def init(self, positions, kernels, ngrid, polmajor=True):
_check(
_bf.bfRomeinInit(self.obj,
asarray(positions).as_BFarray(),
asarray(kernels).as_BFarray(), ngrid, polmajor))
def quantize(src, dst, scale=1.):
src_bf = asarray(src).as_BFarray()
dst_bf = asarray(dst).as_BFarray()
_check(_bf.bfQuantize(src_bf, dst_bf, scale))
return dst
def set_positions(self, positions):
_check(
_bf.bfRomeinSetPositions(self.obj,
asarray(positions).as_BFarray()))
