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()))