Python cupy.get_array_module() Examples

def train_gmm(X, max_iter, tol, means, covariances):
    xp = cupy.get_array_module(X)
    lower_bound = -np.infty
    converged = False
    weights = xp.array([0.5, 0.5], dtype=np.float32)
    inv_cov = 1 / xp.sqrt(covariances)

    for n_iter in range(max_iter):
        prev_lower_bound = lower_bound
        log_prob_norm, log_resp = e_step(X, inv_cov, means, weights)
        weights, means, covariances = m_step(X, xp.exp(log_resp))
        inv_cov = 1 / xp.sqrt(covariances)
        lower_bound = log_prob_norm
        change = lower_bound - prev_lower_bound
        if abs(change) < tol:
            converged = True
            break

    if not converged:
        print('Failed to converge. Increase max-iter or tol.')

    return inv_cov, means, weights, covariances 

def get_array_module(array):
    """Gets an appropriate module from :mod:`numpy` or :mod:`cupy`.

    This is almost equivalent to :func:`cupy.get_array_module`. The differences
    are that this function can be used even if cupy is not available.

    Args:
        array: Input array.

    Returns:
        module: :mod:`cupy` or :mod:`numpy` is returned based on input.
    """
    if config.cupy_enabled:
        return cp.get_array_module(array)
    else:
        return np 

def cfft2(x):
    in_shape = x.shape
    # if both dimensions are odd
    if gpu_config.use_gpu:
        xp = cp.get_array_module(x)
    else:
        xp = np
    if in_shape[0] % 2 == 1 and in_shape[1] % 2 == 1:
        xf = xp.fft.fftshift(xp.fft.fftshift(fft2(x), 0), 1).astype(xp.complex64)
    else:
        out_shape = list(in_shape)
        out_shape[0] =  in_shape[0] + (in_shape[0] + 1) % 2
        out_shape[1] =  in_shape[1] + (in_shape[1] + 1) % 2
        out_shape = tuple(out_shape)
        xf = xp.zeros(out_shape, dtype=xp.complex64)
        xf[:in_shape[0], :in_shape[1]] = xp.fft.fftshift(xp.fft.fftshift(fft2(x), 0), 1).astype(xp.complex64)
        if out_shape[0] != in_shape[0]:
            xf[-1,:] = xp.conj(xf[0,::-1])
        if out_shape[1] != in_shape[1]:
            xf[:,-1] = xp.conj(xf[::-1,0])
    return xf 

def cfft2(x):
    in_shape = x.shape
    # if both dimensions are odd
    if config.use_gpu:
        xp = cp.get_array_module(x)
    else:
        xp = np
    if in_shape[0] % 2 == 1 and in_shape[1] % 2 == 1:
        xf = xp.fft.fftshift(xp.fft.fftshift(fft2(x), 0), 1).astype(xp.complex64)
    else:
        out_shape = list(in_shape)
        out_shape[0] =  in_shape[0] + (in_shape[0] + 1) % 2
        out_shape[1] =  in_shape[1] + (in_shape[1] + 1) % 2
        out_shape = tuple(out_shape)
        xf = xp.zeros(out_shape, dtype=xp.complex64)
        xf[:in_shape[0], :in_shape[1]] = xp.fft.fftshift(xp.fft.fftshift(fft2(x), 0), 1).astype(xp.complex64)
        if out_shape[0] != in_shape[0]:
            xf[-1,:] = xp.conj(xf[0,::-1])
        if out_shape[1] != in_shape[1]:
            xf[:,-1] = xp.conj(xf[::-1,0])
    return xf 

def gradient_loss(generated, truth):
	"""

	:param generated: generated image by the generator at any scale
	:param truth: The ground truth image at that scale
	:return: GDL loss
	"""
	xp = cp.get_array_module(generated.data)
	n, c, h, w = generated.shape
	wx = xp.array([[[1, -1]]]*c, ndmin=4).astype(xp.float32)
	wy = xp.array([[[1], [-1]]]*c, ndmin=4).astype(xp.float32)

	d_gx = F.convolution_2d(generated, wx)
	d_gy = F.convolution_2d(generated, wy)

	d_tx = F.convolution_2d(truth, wx)
	d_ty = F.convolution_2d(truth, wy)

	return (F.sum(F.absolute(d_gx - d_tx)) + F.sum(F.absolute(d_gy - d_ty))) 

def _init_proj_matrix(self, init_sample, compressed_dim, proj_method):
        """
            init the projection matrix
        """
        if gpu_config.use_gpu:
            xp = cp.get_array_module(init_sample[0])
        else:
            xp = np
        x = [xp.reshape(x, (-1, x.shape[2])) for x in init_sample]
        x = [z - z.mean(0) for z in x]
        proj_matrix_ = []
        if self.config.proj_init_method == 'pca':
            for x_, compressed_dim_  in zip(x, compressed_dim):
                proj_matrix, _, _ = xp.linalg.svd(x_.T.dot(x_))
                proj_matrix = proj_matrix[:, :compressed_dim_]
                proj_matrix_.append(proj_matrix)
        elif self.config.proj_init_method == 'rand_uni':
            for x_, compressed_dim_ in zip(x, compressed_dim):
                proj_matrix = xp.random.uniform(size=(x_.shape[1], compressed_dim_))
                proj_matrix /= xp.sqrt(xp.sum(proj_matrix**2, axis=0, keepdims=True))
                proj_matrix_.append(proj_matrix)
        return proj_matrix_ 

def get_array_module(*args):
    """Gets an appropriate one from :mod:`numpy` or :mod:`cupy`.

    This is almost equivalent to :func:`cupy.get_array_module`. The differences
    are that this function can be used even if CUDA is not available and that
    it will return their data arrays' array module for
    :class:`~chainer.Variable` arguments.

    .. deprecated:: v5.0.0

        This API is deprecated. Please use
        :func:`~chainer.backend.get_array_module` instead.

    Args:
        args: Values to determine whether NumPy or CuPy should be used.

    Returns:
        module: :mod:`cupy` or :mod:`numpy` is returned based on the types of
        the arguments.

    """
    return chainer.backend.get_array_module(*args) 

def draw(X, pred, means, covariances, output):
    xp = cupy.get_array_module(X)
    for i in range(2):
        labels = X[pred == i]
        if xp is cupy:
            labels = labels.get()
        plt.scatter(labels[:, 0], labels[:, 1], c=np.random.rand(1, 3))
    if xp is cupy:
        means = means.get()
        covariances = covariances.get()
    plt.scatter(means[:, 0], means[:, 1], s=120, marker='s', facecolors='y',
                edgecolors='k')
    x = np.linspace(-5, 5, 1000)
    y = np.linspace(-5, 5, 1000)
    X, Y = np.meshgrid(x, y)
    for i in range(2):
        dist = stats.multivariate_normal(means[i], covariances[i])
        Z = dist.pdf(np.stack([X, Y], axis=-1))
        plt.contour(X, Y, Z)
    plt.savefig(output) 

def get_array_module(x):
    """Returns the array module for `x`.

    Args:
        x (dezero.Variable or numpy.ndarray or cupy.ndarray): Values to
            determine whether NumPy or CuPy should be used.

    Returns:
        module: `cupy` or `numpy` is returned based on the argument.
    """
    if isinstance(x, Variable):
        x = x.data

    if not gpu_enable:
        return np
    xp = cp.get_array_module(x)
    return xp 

def _init_proj_matrix(self, init_sample, compressed_dim, proj_method):
        """
            init the projection matrix
        """
        if config.use_gpu:
            xp = cp.get_array_module(init_sample[0])
        else:
            xp = np
        x = [xp.reshape(x, (-1, x.shape[2])) for x in init_sample]
        x = [z - z.mean(0) for z in x]
        proj_matrix_ = []
        if config.proj_init_method == 'pca':
            for x_, compressed_dim_  in zip(x, compressed_dim):
                proj_matrix, _, _ = xp.linalg.svd(x_.T.dot(x_))
                proj_matrix = proj_matrix[:, :compressed_dim_]
                proj_matrix_.append(proj_matrix)
        elif config.proj_init_method == 'rand_uni':
            for x_, compressed_dim_ in zip(x, compressed_dim):
                proj_matrix = xp.random.uniform(size=(x_.shape[1], compressed_dim_))
                proj_matrix /= xp.sqrt(xp.sum(proj_matrix**2, axis=0, keepdims=True))
                proj_matrix_.append(proj_matrix)
        return proj_matrix_ 

def circ_mul_v(circ,v,eigs=None):
    ''' multiply a circulant matrix A by multi vector v.

    Args:
        circ (ndarray): representation of the multilevel circulant matrix A, i.e.
             the first column of A in proper shape.
        v (ndarray): vector to be multiplied. Should be reshaped to the same shape
             as circ. Should be the same reshape order (column first/row first) as circ.
    Returns:
        result of multiplication.
    '''
    if use_gpu > 0:
        import cupy
        xp = cupy.get_array_module(circ)
    else:
        xp = np
    
    if eigs is None:
        eigs = circ_eigs(circ)
    tmp = xp.real(xp.fft.ifft2(xp.fft.fft2(v,norm='ortho')*eigs,norm='ortho'))
    if xp is cupy:
        return tmp.astype(xp.float32)
    else:
        return tmp 

def fit(A, b, tol, max_iter):
    # Note that this function works even tensors 'A' and 'b' are NumPy or CuPy
    # arrays.
    xp = cupy.get_array_module(A)
    x = xp.zeros_like(b, dtype=np.float64)
    r0 = b - xp.dot(A, x)
    p = r0
    for i in range(max_iter):
        a = xp.inner(r0, r0) / xp.inner(p, xp.dot(A, p))
        x += a * p
        r1 = r0 - a * xp.dot(A, p)
        if xp.linalg.norm(r1) < tol:
            return x
        b = xp.inner(r1, r1) / xp.inner(r0, r0)
        p = r1 + b * p
        r0 = r1
    print('Failed to converge. Increase max-iter or tol.')
    return x 

def symmetrize_filter(hf):
    """
        ensure hermetian symmetry
    """
    if gpu_config.use_gpu:
        xp = cp.get_array_module(hf[0])
    else:
        xp = np
    for i in range(len(hf)):
        dc_ind = int((hf[i].shape[0]+1) / 2)
        hf[i][dc_ind:, -1, :] = xp.conj(xp.flipud(hf[i][:dc_ind-1, -1, :]))
    return hf 

def _find_gram_vector(self, samplesf, new_sample, num_training_samples):
        if gpu_config.use_gpu:
            xp = cp.get_array_module(samplesf[0])
        else:
            xp = np
        gram_vector = xp.inf * xp.ones((self.config.num_samples))
        if num_training_samples > 0:
            ip = 0.
            for k in range(len(new_sample)):
                samplesf_ = samplesf[k][:, :, :, :num_training_samples]
                samplesf_ = samplesf_.reshape((-1, num_training_samples))
                new_sample_ = new_sample[k].flatten()
                ip += xp.real(2 * samplesf_.T.dot(xp.conj(new_sample_)))
            gram_vector[:num_training_samples] = ip
        return gram_vector 

def get_array_module(*args, **kwargs):
    """ default array module when cupy is not present """
    return np 

def shift_sample(xf, shift, kx, ky):
    if gpu_config.use_gpu:
        xp = cp.get_array_module(xf[0])
    else:
        xp = np
    shift_exp_y = [xp.exp(1j * shift[0] * ky_).astype(xp.complex64) for ky_ in ky]
    shift_exp_x = [xp.exp(1j * shift[1] * kx_).astype(xp.complex64) for kx_ in kx]
    xf = [xf_ * sy_.reshape(-1, 1, 1, 1) * sx_.reshape((1, -1, 1, 1))
            for xf_, sy_, sx_ in zip(xf, shift_exp_y, shift_exp_x)]
    return xf 

def full_fourier_coeff(xf):
    """
        Reconstructs the full Fourier series coefficients
    """
    if gpu_config.use_gpu:
        xp = cp.get_array_module(xf[0])
    else:
        xp = np
    xf = [xp.concatenate([xf_, xp.conj(xp.rot90(xf_[:, :-1,:], 2))], axis=1) for xf_ in xf]
    return xf 

def cifft2(xf):
    if gpu_config.use_gpu:
        xp = cp.get_array_module(xf)
    else:
        xp = np
    x = xp.real(ifft2(xp.fft.ifftshift(xp.fft.ifftshift(xf, 0),1))).astype(xp.float32)
    return x 

def ifft2(x):
    if gpu_config.use_gpu:
        xp = cp.get_array_module(x)
    else:
        xp = np
    return xp.fft.ifft(xp.fft.ifft(x, axis=1), axis=0).astype(xp.complex64)
    # return ifft(ifft(x, axis=1), axis=0) 

def _proj_sample(self, x, P):
        if gpu_config.use_gpu:
            xp = cp.get_array_module(x[0])
        else:
            xp = np
        return [xp.matmul(P_.T, x_) for x_, P_ in zip(x, P)] 

def inner_product_joint(xf, yf):
    """
        computes the joint inner product between two filters and projection matrices
    """
    if gpu_config.use_gpu:
        xp = cp.get_array_module(xf[0][0])
    else:
        xp = np
    ip = 0
    for i in range(len(xf[0])):
        ip += 2 * xp.vdot(xf[0][i].flatten(), yf[0][i].flatten()) - xp.vdot(xf[0][i][:, -1, :].flatten(), yf[0][i][:, -1, :].flatten())
        ip += xp.vdot(xf[1][i].flatten(), yf[1][i].flatten())
    return xp.real(ip) 

def train_filter(hf, samplesf, yf, reg_filter, sample_weights, sample_energy, reg_energy, CG_opts, CG_state,config):
    """
        do conjugate graident optimization of the filter
    """
    if gpu_config.use_gpu:
        xp = cp.get_array_module(hf[0][0])
    else:
        xp = np
    # construct the right hand side vector (A^H weight yf)
    rhs_samplef = [xp.matmul(xf, sample_weights) for xf in samplesf]
    rhs_samplef = [(xp.conj(xf) * yf[:,:,xp.newaxis,xp.newaxis])
            for xf, yf in zip(rhs_samplef, yf)]

    # construct preconditioner
    diag_M = [(1 - config.precond_reg_param) * (config.precond_data_param * m + (1 - config.precond_data_param) * xp.mean(m, 2, keepdims=True)) + \
              config.precond_reg_param * reg_energy_ for m, reg_energy_ in zip(sample_energy, reg_energy)]
    hf, _, CG_state = preconditioned_conjugate_gradient(
            lambda x: lhs_operation(x, samplesf, reg_filter, sample_weights), # A
            [rhs_samplef],                                                    # b
            CG_opts,                                                          # opts
            lambda x: diag_precond(x, [diag_M]),                              # M1
            None,                                                             # M2
            inner_product_filter,
            [hf],
            CG_state)
    # res_norms = res_norms / xp.sqrt(inner_product_filter([rhs_samplef], [rhs_samplef]))
    return hf[0], CG_state #res_norms, CG_state 

def test_get_array_module(self, xp, dtype):
        def func(x):
            assert xp == cupy.get_array_module(x)
            return x

        return func 

def block_toep2_sym(a):
    '''generate full representation of 2-level symmetric toeplitz matrix

    Args:
        a (ndarray): 1-st column of the symmetrec toeplitz matrix in proper shape.
    Returns:
        Full filled toeplitz matrix.
    '''
    if use_gpu > 0:
        import cupy
        xp = cupy.get_array_module(a)
        if xp is cupy:
            a = xp.asnumpy(a)
    else:
        xp = np
        a = np.asnumpy(a)
        
    A1 = []
    n0,n1 = a.shape
    for i in range(n1):
        A1.append(splin.toeplitz(a[:,i]))
    A = np.empty((n1,n0,n1,n0))
    for i in range(n1):
        for j in range(n1):
            A[i,:,j,:] = A1[np.int(np.abs(i-j))]
    A.shape = (n0*n1,n0*n1)
    A = xp.asarray(A)
    
    return(A) 

def circ_eigs(circ,dtype=np.complex64):
    ''' calculate eigenvalues of multilevel circulant matrix A

    Args:
        circ (ndarray): representation of the multilevel circulant matrix A, i.e.
             the first column of A in proper shape.
    Returns:
        eigenvalues of A.
    '''
    if use_gpu > 0:
        import cupy
        xp = cupy.get_array_module(circ)
    else:
        xp = np
    return xp.fft.fft2(circ,norm='ortho').astype(dtype)*xp.sqrt(np.prod(circ.shape)) 

def embed_toep2circ(toep,v=None):
    '''embed toeplitz matrix to circulant matrix.

    Args:
        toep (ndarray): representation of multilevel toeplitz matrix, i.e.
            the first column of A in proper shape.
        v (ndarray): embed a vector together with toep.
    Returns:
        representation of embedded multilevel circulant matrix and embedded vector.
    '''
    if use_gpu > 0:
        import cupy
        xp = cupy.get_array_module(toep)
    else:
        xp = np
        
#    circ = xp.zeros((2*xp.array(toep.shape)).astype(xp.int))
    circ = xp.zeros((2*toep.shape[0],2*toep.shape[1]))
    s = []
    for idim in range(toep.ndim):
        s.append(slice(0,toep.shape[idim]))
    circ[tuple(s)] = toep
    if not v is None:
        if v.ndim == toep.ndim:
            resv = xp.zeros_like(circ)
            resv[tuple(s)] = v
        else:
            resv = xp.zeros((v.shape[0],*circ.shape))
            resv[tuple(slice(None),*s)] = v
    for idim in range(toep.ndim):
        s[idim] = slice(toep.shape[idim]+1,None)
        s2 = s.copy()
        s2[idim] = slice(toep.shape[idim]-1,0,-1)
        circ[tuple(s)] = circ[tuple(s2)]
        s[idim] = slice(None)
    if v is None:
        return circ
    else:
        return circ,resv 

def __init__(self,toep,*args,**kwargs):
        '''Input toep is the first row of the underlying Toeplitz matrix'''
        if use_gpu > 0:
            import cupy
            self.xp = cupy.get_array_module(toep)
        else:
            self.xp = np
            
        self.nz,self.ny,self.nx = toep.shape
        self.eigs = self.xp.zeros((self.nz,2*self.ny,2*self.nx),dtype=np.complex64)
        for i in range(self.nz):
            tmptoep = embed_toep2circ(toep[i])
            self.eigs[i] = circ_eigs(tmptoep,dtype=np.complex64)
        self.dtype = toep.dtype
        self.get_m_eigs(toep) 

def cg(A, b, x=None, tol=1.0e-5, max_iter=None):
    # Note that this function works even tensors 'A' and 'b' are NumPy or CuPy
    # arrays.
    if use_gpu > 0:
        import cupy
        xp = cupy.get_array_module(b)
    else:
        xp = np
    if max_iter is None:
        max_iter = 10*len(b)
    if x is None:
        x = xp.zeros_like(b, dtype=np.float32)
    r0 = b - A.matvec(x)
    p = r0
    now = time.time()
    for i in range(max_iter):
        a = xp.inner(r0, r0) / xp.inner(p, A.matvec(p))
        x += a * p
        r1 = r0 - a * A.matvec(p)
        res = xp.linalg.norm(r1)
        if  res < tol:
            return x
        b = xp.inner(r1, r1) / xp.inner(r0, r0)
        p = r1 + b * p
        r0 = r1
        if time.time() - now > 10:
            now = time.time()
            print('iter: {:8d}  residual: {:16.7e}'.format(i,xp.asnumpy(res)))
            #print('iter: {:8d}  residual: {:16.7e}'.format(i,cupy.asnumpy(res)))
    print('Failed to converge. Increase max-iter or tol.')
    return x 

def VGGprepare_am_input(var):
	xp = get_array_module(var)

	# var = F.resize_images(var, size)
	var = F.transpose(var, (0, 2, 3, 1)) # [[W, H, C]]
	var = F.flip(var, 3)
	var -= xp.array([[103.939, 116.779, 123.68]], dtype=xp.float32)
	var = F.transpose(var, (0, 3, 1, 2))
	return var 

def predict(self, x, no_of_predictions=1, seq_len=4):
		"""

		:param x:
		:param no_of_predictions:
		:param seq_len:
		:return:
		"""
		# x shape = [n, 12, h, w]
		xp = cp.get_array_module(x)
		n, c, h, w = x.shape
		outputs = []

		for i in range(no_of_predictions):
			print("Predicting frame no : ",i+1)
			seq = resize_images(x, (int(h / 2 ** 3), int(w / 2 ** 3)))
			print((int(h / 2 ** 3), int(w / 2 ** 3)))
			output = None

			for j in range(1, 5):
				output = self.singleforward(j, seq, output)
				if j != 4:
					seq = resize_images(x, (int(h / 2 ** (3-j)), int(w / 2 ** (3-j))))

			outputs.append(output.data)
			x = xp.concatenate([x, output.data], 1)[:, -seq_len*3:, :, :]
			print("Predictions done for : ", i+1)
		return outputs