def blk_matvec(x, blks): nargins = [[blk.shape[-1] for blk in blkrow] for blkrow in blks] nargouts = [[blk.shape[0] for blk in blkrow] for blkrow in blks] nargin = sum(nargins[0]) nargout = sum([out[0] for out in nargouts]) nx = len(x) self.logger.debug('Multiplying with a vector of size %d' % nx) self.logger.debug('nargin=%d, nargout=%d' % (nargin, nargout)) if nx != nargin: raise ShapeError('Multiplying with vector of wrong shape.') result_type = np.result_type(self.dtype, x.dtype) y = np.zeros(nargout, dtype=result_type) nblk_row = len(blks) nblk_col = len(blks[0]) row_start = col_start = 0 for row in range(nblk_row): row_end = row_start + nargouts[row][0] yout = y[row_start:row_end] for col in range(nblk_col): col_end = col_start + nargins[0][col] xin = x[col_start:col_end] B = blks[row][col] yout[:] += B * xin col_start = col_end row_start = row_end col_start = 0 return y
def blk_matvec(x, blks): nx = len(x) nargins = [blk.shape[-1] for blk in blocks] nargin = sum(nargins) nargouts = [blk.shape[0] for blk in blocks] nargout = sum(nargouts) self.logger.debug('Multiplying with a vector of size %d' % nx) self.logger.debug('nargin=%d, nargout=%d' % (nargin, nargout)) if nx != nargin: raise ShapeError('Multiplying with vector of wrong shape.') result_type = np.result_type(self.dtype, x.dtype) y = np.empty(nargout, dtype=result_type) nblks = len(blks) row_start = col_start = 0 for blk in range(nblks): row_end = row_start + nargouts[blk] yout = y[row_start:row_end] col_end = col_start + nargins[blk] xin = x[col_start:col_end] B = blks[blk] yout[:] = B * xin col_start = col_end row_start = row_end return y
def update_block(self, row, column, new_block): if new_block.shape != self._blocks[row][column].shape: raise ShapeError( 'The new block should have the same shape as the block you are trying to replace' ) self._blocks[row][column] = new_block if self.symmetric: self._blocks[column][row] = new_block.T return
def __init__(self, blocks, symmetric=False, **kwargs): # If building a symmetric operator, fill in the blanks. # They're just references to existing objects. if symmetric: nrow = len(blocks) ncol = len(blocks[0]) if nrow != ncol: raise ShapeError('Inconsistent shape.') for block_row in blocks: if not block_row[0].symmetric: raise ValueError('Blocks on diagonal must be symmetric.') self._blocks = blocks[:] for i in range(1, nrow): for j in range(i - 1, -1, -1): self._blocks[i].insert(0, self._blocks[j][i].T) else: self._blocks = blocks log = kwargs.get('logger', null_log) log.debug('Building new BlockLinearOperator') nargins = [[blk.shape[-1] for blk in row] for row in self._blocks] log.debug('nargins = ' + repr(nargins)) nargins_by_row = [nargin[0] for nargin in nargins] if min(nargins_by_row) != max(nargins_by_row): raise ShapeError('Inconsistent block shapes') nargouts = [[blk.shape[0] for blk in row] for row in self._blocks] log.debug('nargouts = ' + repr(nargouts)) for row in nargouts: if min(row) != max(row): raise ShapeError('Inconsistent block shapes') nargin = sum(nargins[0]) nargout = sum([out[0] for out in nargouts]) # Create blocks of transpose operator. blocksT = map(lambda *row: [blk.T for blk in row], *self._blocks) def blk_matvec(x, blks): nargins = [[blk.shape[-1] for blk in blkrow] for blkrow in blks] nargouts = [[blk.shape[0] for blk in blkrow] for blkrow in blks] nargin = sum(nargins[0]) nargout = sum([out[0] for out in nargouts]) nx = len(x) self.logger.debug('Multiplying with a vector of size %d' % nx) self.logger.debug('nargin=%d, nargout=%d' % (nargin, nargout)) if nx != nargin: raise ShapeError('Multiplying with vector of wrong shape.') result_type = np.result_type(self.dtype, x.dtype) y = np.zeros(nargout, dtype=result_type) nblk_row = len(blks) nblk_col = len(blks[0]) row_start = col_start = 0 for row in range(nblk_row): row_end = row_start + nargouts[row][0] yout = y[row_start:row_end] for col in range(nblk_col): col_end = col_start + nargins[0][col] xin = x[col_start:col_end] B = blks[row][col] yout[:] += B * xin col_start = col_end row_start = row_end col_start = 0 return y flat_blocks = list(itertools.chain(*blocks)) blk_dtypes = [blk.dtype for blk in flat_blocks] op_dtype = np.result_type(*blk_dtypes) super(BlockLinearOperator, self).__init__(nargin, nargout, symmetric=symmetric, matvec=lambda x: blk_matvec(x, self._blocks), matvec_transp=lambda x: blk_matvec(x, blocksT), dtype=op_dtype) self.T._blocks = blocksT