Python numpy.diag_indices() Examples

def energy_nuc(mol):
    atom_charges = mol.atom_charges()
    atom_coords = mol.atom_coords()

    distances = numpy.linalg.norm(atom_coords[:,None,:] - atom_coords, axis=2)
    distances_in_AA = distances * lib.param.BOHR
    # numerically exclude atomic self-interaction terms
    distances_in_AA[numpy.diag_indices(mol.natm)] = 1e60

    # one atom is H, another atom is N or O
    where_NO = (atom_charges == 7) | (atom_charges == 8)
    mask = (atom_charges[:,None] == 1) & where_NO
    mask = mask | mask.T
    scale = alpha = _get_alpha(atom_charges[:,None], atom_charges)
    scale[mask] *= numpy.exp(-distances_in_AA[mask])
    scale[~mask] = numpy.exp(-alpha[~mask] * distances_in_AA[~mask])

    gamma = _get_gamma(mol)
    z_eff = mopac_param.CORE[atom_charges]
    e_nuc = .5 * numpy.einsum('i,ij,j->', z_eff, gamma, z_eff)
    e_nuc += .5 * numpy.einsum('i,j,ij,ij->', z_eff, z_eff, scale,
                               mopac_param.E2/distances_in_AA - gamma)
    return e_nuc 

def test_field_gradients(self):
        mol = gto.M(atom='H1 0.5 -0.6 0.4; H2 -0.5, 0.4, -0.3; H -0.4 -0.3 0.5; H 0.3 0.5 -0.6',
                    unit='B',
                    basis={'H': [[0,[2., 1]]], 'H1':[[1,[.5, 1]]], 'H2':[[1,[1,1]]]})
        grids = dft.gen_grid.Grids(mol)
        grids.build()
        ao = mol.eval_gto('GTOval', grids.coords)
        r0 = mol.atom_coord(0)
        dr = grids.coords - r0
        dd = numpy.linalg.norm(dr, axis=1)
        rr = 3 * numpy.einsum('ix,iy->ixy', dr, dr)
        for i in range(3):
            rr[:,i,i] -= dd**2
        h1ref = lib.einsum('i,ixy,ip,iq->xypq', grids.weights/dd**5, rr, ao, ao)

        h1ao = rhf_ssc._get_integrals_fcsd(mol, 0)
        h1ao[numpy.diag_indices(3)] += rhf_ssc._get_integrals_fc(mol, 0)
        self.assertAlmostEqual(abs(h1ref - h1ao).max(), 0, 5) 

def dist_diff(self, other=None):
        r"""
        Calculate distances and vectors between atoms.

        Args:
            other (:class:`~berny.Geometry`): calculate distances between two
                geometries if given or within a geometry if not

        Returns:
            :math:`R_{ij}:=|\mathbf R_i-\mathbf R_j|` and
            :math:`R_{ij\alpha}:=(\mathbf R_i)_\alpha-(\mathbf R_j)_\alpha`.
        """
        if other is None:
            other = self
        diff = self.coords[:, None, :] - other.coords[None, :, :]
        dist = np.sqrt(np.sum(diff ** 2, 2))
        dist[np.diag_indices(len(self))] = np.inf
        return dist, diff 

def __init__(self, features, using_cache=False, identity_init=True, eps=1e-3):
        super().__init__(features, using_cache)

        self.eps = eps

        self.lower_indices = np.tril_indices(features, k=-1)
        self.upper_indices = np.triu_indices(features, k=1)
        self.diag_indices = np.diag_indices(features)

        n_triangular_entries = ((features - 1) * features) // 2

        self.lower_entries = nn.Parameter(torch.zeros(n_triangular_entries))
        self.upper_entries = nn.Parameter(torch.zeros(n_triangular_entries))
        self.unconstrained_upper_diag = nn.Parameter(torch.zeros(features))

        self._initialize(identity_init) 

def _get_knmat(exog, xcov, sl):
    # Utility function, see equation 2.2 of Barber & Candes.

    nobs, nvar = exog.shape

    ash = np.linalg.inv(xcov)
    ash *= -np.outer(sl, sl)
    i, j = np.diag_indices(nvar)
    ash[i, j] += 2 * sl

    umat = np.random.normal(size=(nobs, nvar))
    u, _ = np.linalg.qr(exog)
    umat -= np.dot(u, np.dot(u.T, umat))
    umat, _ = np.linalg.qr(umat)

    ashr, xc, _ = np.linalg.svd(ash, 0)
    ashr *= np.sqrt(xc)
    ashr = ashr.T

    ex = (sl[:, None] * np.linalg.solve(xcov, exog.T)).T
    exogn = exog - ex + np.dot(umat, ashr)

    return exogn 

def test_get_data_dst_nodata(rast, dst_nodata, dst_arr):
    fp = rast.fp.dilate(1)
    inner_slice = rast.fp.slice_in(fp)
    rast.set_data(dst_arr, channels=None)

    arr = rast.get_data(band=[-1], dst_nodata=dst_nodata, fp=fp)

    arr2 = rast.get_data(band=[-1], dst_nodata=dst_nodata, fp=fp.erode(1))
    assert np.all(arr2 == arr[inner_slice])

    if rast.nodata is not None:
        assert np.all(arr[inner_slice][np.diag_indices(dst_arr.shape[0])] == dst_nodata)
        arr[inner_slice][np.diag_indices(dst_arr.shape[0])] = rast.nodata
    assert np.all(arr[inner_slice] == dst_arr)

    outer_mask = np.ones(fp.shape, bool)
    outer_mask[inner_slice] = False
    assert np.all(arr[outer_mask] == dst_nodata) 

def compute_outflow_pdf(self):
        """
        Lifetime model. The method compute outflow_pdf returns an array year-by-cohort of the probability of a item added to stock in year m (aka cohort m) leaves in in year n. This value equals pdf(n,m).
        This is the only method for the inflow-driven model where the lifetime distribution directly enters the computation. All other stock variables are determined by mass balance.
        The shape of the output pdf array is NoofYears * NoofYears, but the meaning is years by age-cohorts.
        The method does nothing if the pdf alreay exists.
        """
        if self.pdf is None:
            self.compute_sf() # computation of pdfs moved to this method: compute survival functions sf first, then calculate pdfs from sf.
            self.pdf   = np.zeros((len(self.t), len(self.t)))
            self.pdf[np.diag_indices(len(self.t))] = np.ones(len(self.t)) - self.sf.diagonal(0)
            for m in range(0,len(self.t)):
                self.pdf[np.arange(m+1,len(self.t)),m] = -1 * np.diff(self.sf[np.arange(m,len(self.t)),m])            
            return self.pdf
        else:
            # pdf already exists
            return self.pdf 

def _initialize_error_cov_diagonal(self, scalar=False):
        # Initialize the parameters
        self.parameters['error_cov'] = 1 if scalar else self.k_endog

        # Setup fixed components of state space matrices

        # Setup indices of state space matrices
        k_endog = self.k_endog
        k_factors = self.k_factors
        idx = np.diag_indices(k_endog)
        if self.error_order > 0:
            matrix = 'state_cov'
            idx = (idx[0] + k_factors, idx[1] + k_factors)
        else:
            matrix = 'obs_cov'
        self._idx_error_cov = (matrix,) + idx 

def _get_knmat(exog, xcov, sl):
    # Utility function, see equation 2.2 of Barber & Candes.

    nobs, nvar = exog.shape

    ash = np.linalg.inv(xcov)
    ash *= -np.outer(sl, sl)
    i, j = np.diag_indices(nvar)
    ash[i, j] += 2 * sl

    umat = np.random.normal(size=(nobs, nvar))
    u, _ = np.linalg.qr(exog)
    umat -= np.dot(u, np.dot(u.T, umat))
    umat, _ = np.linalg.qr(umat)

    ashr, xc, _ = np.linalg.svd(ash, 0)
    ashr *= np.sqrt(xc)
    ashr = ashr.T

    ex = (sl[:, None] * np.linalg.solve(xcov, exog.T)).T
    exogn = exog - ex + np.dot(umat, ashr)

    return exogn 

def __init__(self, one_body, two_body, constant=0.):
        if two_body.dtype != numpy.float:
            raise ValueError('Two-body tensor has invalid dtype. Expected {} '
                             'but was {}'.format(numpy.float, two_body.dtype))
        if not numpy.allclose(two_body, two_body.T):
            raise ValueError('Two-body tensor must be symmetric.')
        if not numpy.allclose(one_body, one_body.T.conj()):
            raise ValueError('One-body tensor must be Hermitian.')

        # Move the diagonal of two_body to one_body
        diag_indices = numpy.diag_indices(one_body.shape[0])
        one_body[diag_indices] += two_body[diag_indices]
        numpy.fill_diagonal(two_body, 0.)

        self.one_body = one_body
        self.two_body = two_body
        self.constant = constant 

def test_antidiagonal(self):
        m, n = (3, 3)
        Q = numpy.zeros((m, n), dtype=complex)
        Q[0, 2] = 1.
        Q[1, 1] = 1.
        Q[2, 0] = 1.
        givens_rotations, V, diagonal = givens_decomposition(Q)

        # There should be no Givens rotations
        self.assertEqual(givens_rotations, list())

        # VQ should equal the diagonal
        VQ = V.dot(Q)
        D = numpy.zeros((m, n), dtype=complex)
        D[numpy.diag_indices(m)] = diagonal
        for i in range(n):
            for j in range(n):
                self.assertAlmostEqual(VQ[i, j], D[i, j]) 

def test_square(self):
        m, n = (3, 3)

        # Obtain a random matrix of orthonormal rows
        Q = random_unitary_matrix(n)
        Q = Q[:m, :]
        Q = Q[:m, :]

        # Get Givens decomposition of Q
        givens_rotations, V, diagonal = givens_decomposition(Q)

        # There should be no Givens rotations
        self.assertEqual(givens_rotations, list())

        # Compute V * Q * U^\dagger
        W = V.dot(Q)

        # Construct the diagonal matrix
        D = numpy.zeros((m, n), dtype=complex)
        D[numpy.diag_indices(m)] = diagonal

        # Assert that W and D are the same
        for i in range(m):
            for j in range(n):
                self.assertAlmostEqual(D[i, j], W[i, j]) 

def diag_indices_from(arr):
    """
    Return the indices to access the main diagonal of an n-dimensional array.
    See `diag_indices` for full details.

    Args:
        arr (cupy.ndarray): At least 2-D.

     .. seealso:: :func:`numpy.diag_indices_from`

    """
    if not isinstance(arr, cupy.ndarray):
        raise TypeError("Argument must be cupy.ndarray")

    if not arr.ndim >= 2:
        raise ValueError("input array must be at least 2-d")
    # For more than d=2, the strided formula is only valid for arrays with
    # all dimensions equal, so we check first.
    if not cupy.all(cupy.diff(arr.shape) == 0):
        raise ValueError("All dimensions of input must be of equal length")

    return diag_indices(arr.shape[0], arr.ndim) 

def test_orth(self):
        numpy.random.seed(10)
        n = 100
        a = numpy.random.random((n,n))
        s = numpy.dot(a.T, a)
        c = orth.lowdin(s)
        self.assertTrue(numpy.allclose(reduce(numpy.dot, (c.T, s, c)),
                                       numpy.eye(n)))
        x1 = numpy.dot(a, c)
        x2 = orth.vec_lowdin(a)
        d = numpy.dot(x1.T,x2)
        d[numpy.diag_indices(n)] = 0
        self.assertAlmostEqual(numpy.linalg.norm(d), 0, 9)
        self.assertAlmostEqual(numpy.linalg.norm(c), 36.56738258719514, 9)
        self.assertAlmostEqual(abs(c).sum(), 2655.5580057303964, 7) 

def test_schmidt(self):
        numpy.random.seed(10)
        n = 100
        a = numpy.random.random((n,n))
        s = numpy.dot(a.T, a)
        c = orth.schmidt(s)
        self.assertTrue(numpy.allclose(reduce(numpy.dot, (c.T, s, c)),
                                       numpy.eye(n)))
        x1 = numpy.dot(a, c)
        x2 = orth.vec_schmidt(a)
        d = numpy.dot(x1.T,x2)
        d[numpy.diag_indices(n)] = 0
        self.assertAlmostEqual(numpy.linalg.norm(d), 0, 9)
        self.assertAlmostEqual(numpy.linalg.norm(c), 36.56738258719514, 9)
        self.assertAlmostEqual(abs(c).sum(), 1123.2089785000373, 7) 

def _from_rhf_init_dm(dm, breaksym=True):
    dma = dm * .5
    dm = scipy.linalg.block_diag(dma, dma)
    if breaksym:
        nao = dma.shape[0]
        idx, idy = numpy.diag_indices(nao)
        dm[idx+nao,idy] = dm[idx,idy+nao] = dma.diagonal() * .05
    return dm 

def _unpack_t2_tril(t2tril, nocc, nvir, out=None, t2sym='jiba'):
    t2 = numpy.ndarray((nocc,nocc,nvir,nvir), dtype=t2tril.dtype, buffer=out)
    idx,idy = numpy.tril_indices(nocc)
    if t2sym == 'jiba':
        t2[idy,idx] = t2tril.transpose(0,2,1)
        t2[idx,idy] = t2tril
    elif t2sym == '-jiba':
        t2[idy,idx] = -t2tril.transpose(0,2,1)
        t2[idx,idy] = t2tril
    elif t2sym == '-jiab':
        t2[idy,idx] =-t2tril
        t2[idx,idy] = t2tril
        t2[numpy.diag_indices(nocc)] = 0
    return t2 

def _gamma1_intermediates(mycc, t1, t2, l1, l2, eris=None):
    doo, dov, dvo, dvv = gccsd_rdm._gamma1_intermediates(mycc, t1, t2, l1, l2)

    if eris is None: eris = mycc.ao2mo()

    nocc, nvir = t1.shape
    bcei = numpy.asarray(eris.ovvv).conj().transpose(3,2,1,0)
    majk = numpy.asarray(eris.ooov).conj().transpose(2,3,0,1)
    bcjk = numpy.asarray(eris.oovv).conj().transpose(2,3,0,1)

    mo_e = eris.mo_energy
    eia = mo_e[:nocc,None] - mo_e[nocc:]
    d3 = lib.direct_sum('ia+jb+kc->ijkabc', eia, eia, eia)

    t3c =(numpy.einsum('jkae,bcei->ijkabc', t2, bcei)
        - numpy.einsum('imbc,majk->ijkabc', t2, majk))
    t3c = t3c - t3c.transpose(0,1,2,4,3,5) - t3c.transpose(0,1,2,5,4,3)
    t3c = t3c - t3c.transpose(1,0,2,3,4,5) - t3c.transpose(2,1,0,3,4,5)
    t3c /= d3

    t3d = numpy.einsum('ia,bcjk->ijkabc', t1, bcjk)
    t3d += numpy.einsum('ai,jkbc->ijkabc', eris.fock[nocc:,:nocc], t2)
    t3d = t3d - t3d.transpose(0,1,2,4,3,5) - t3d.transpose(0,1,2,5,4,3)
    t3d = t3d - t3d.transpose(1,0,2,3,4,5) - t3d.transpose(2,1,0,3,4,5)
    t3d /= d3

    goo = numpy.einsum('iklabc,jklabc->ij', (t3c+t3d).conj(), t3c) * (1./12)
    gvv = numpy.einsum('ijkacd,ijkbcd->ab', t3c+t3d, t3c.conj()) * (1./12)
    doo[numpy.diag_indices(nocc)] -= goo.diagonal()
    dvv[numpy.diag_indices(nvir)] += gvv.diagonal()
    dvo += numpy.einsum('ijab,ijkabc->ck', t2.conj(), t3c) * (1./4)

    return doo, dov, dvo, dvv

# gamma2 intermediates in Chemist's notation 

def _make_rdm1(mycc, d1, with_frozen=True, ao_repr=False):
    r'''dm1[p,q] = <q_alpha^\dagger p_alpha> + <q_beta^\dagger p_beta>

    The convention of 1-pdm is based on McWeeney's book, Eq (5.4.20).
    The contraction between 1-particle Hamiltonian and rdm1 is
    E = einsum('pq,qp', h1, rdm1)
    '''
    doo, dov, dvo, dvv = d1
    nocc, nvir = dov.shape
    nmo = nocc + nvir
    dm1 = numpy.empty((nmo,nmo), dtype=doo.dtype)
    dm1[:nocc,:nocc] = doo + doo.conj().T
    dm1[:nocc,nocc:] = dov + dvo.conj().T
    dm1[nocc:,:nocc] = dm1[:nocc,nocc:].conj().T
    dm1[nocc:,nocc:] = dvv + dvv.conj().T
    dm1[numpy.diag_indices(nocc)] += 2

    if with_frozen and mycc.frozen is not None:
        nmo = mycc.mo_occ.size
        nocc = numpy.count_nonzero(mycc.mo_occ > 0)
        rdm1 = numpy.zeros((nmo,nmo), dtype=dm1.dtype)
        rdm1[numpy.diag_indices(nocc)] = 2
        moidx = numpy.where(mycc.get_frozen_mask())[0]
        rdm1[moidx[:,None],moidx] = dm1
        dm1 = rdm1

    if ao_repr:
        mo = mycc.mo_coeff
        dm1 = lib.einsum('pi,ij,qj->pq', mo, dm1, mo.conj())
    return dm1

# Note vvvv part of 2pdm have been symmetrized.  It does not correspond to
# vvvv part of CI 2pdm 

def _make_rdm1(mycc, d1, with_frozen=True, ao_repr=False):
    r'''
    One-particle density matrix in the molecular spin-orbital representation
    (the occupied-virtual blocks from the orbital response contribution are
    not included).

    dm1[p,q] = <q^\dagger p>  (p,q are spin-orbitals)

    The convention of 1-pdm is based on McWeeney's book, Eq (5.4.20).
    The contraction between 1-particle Hamiltonian and rdm1 is
    E = einsum('pq,qp', h1, rdm1)
    '''
    doo, dov, dvo, dvv = d1
    nocc, nvir = dov.shape
    nmo = nocc + nvir

    dm1 = numpy.empty((nmo,nmo), dtype=doo.dtype)
    dm1[:nocc,:nocc] = doo + doo.conj().T
    dm1[:nocc,nocc:] = dov + dvo.conj().T
    dm1[nocc:,:nocc] = dm1[:nocc,nocc:].conj().T
    dm1[nocc:,nocc:] = dvv + dvv.conj().T
    dm1 *= .5
    dm1[numpy.diag_indices(nocc)] += 1

    if with_frozen and mycc.frozen is not None:
        nmo = mycc.mo_occ.size
        nocc = numpy.count_nonzero(mycc.mo_occ > 0)
        rdm1 = numpy.zeros((nmo,nmo), dtype=dm1.dtype)
        rdm1[numpy.diag_indices(nocc)] = 1
        moidx = numpy.where(mycc.get_frozen_mask())[0]
        rdm1[moidx[:,None],moidx] = dm1
        dm1 = rdm1

    if ao_repr:
        mo = mycc.mo_coeff
        dm1 = lib.einsum('pi,ij,qj->pq', mo, dm1, mo.conj())
    return dm1 

def get_hcore(mol):
    assert(not mol.has_ecp())
    nao = mol.nao

    atom_charges = mol.atom_charges()
    basis_atom_charges = atom_charges[mol._bas[:,gto.ATOM_OF]]

    basis_ip = []
    basis_u = []
    for i, z in enumerate(basis_atom_charges):
        l = mol.bas_angular(i)
        if l == 0:
            basis_ip.append(mopac_param.VS[z])
            basis_u.append(mopac_param.USS3[z])
        else:
            basis_ip.append(mopac_param.VP[z])
            basis_u.append(mopac_param.UPP3[z])

    ao_atom_charges = _to_ao_labels(mol, basis_atom_charges)
    ao_ip = _to_ao_labels(mol, basis_ip)

    # Off-diagonal terms
    hcore  = mol.intor('int1e_ovlp')
    hcore *= ao_ip[:,None] + ao_ip
    hcore *= _get_beta0(ao_atom_charges[:,None], ao_atom_charges)

    # U term 
    hcore[numpy.diag_indices(nao)] = _to_ao_labels(mol, basis_u)

    # Nuclear attraction
    gamma = _get_gamma(mol)
    z_eff = mopac_param.CORE[atom_charges]
    vnuc = numpy.einsum('ij,j->i', gamma, z_eff)

    aoslices = mol.aoslice_by_atom()
    for ia, (p0, p1) in enumerate(aoslices[:,2:]):
        idx = numpy.arange(p0, p1)
        hcore[idx,idx] -= vnuc[ia]
    return hcore 

def _get_gamma(mol, F03=mopac_param.F03):
    atom_charges = mol.atom_charges()
    atom_coords = mol.atom_coords()
    distances = numpy.linalg.norm(atom_coords[:,None,:] - atom_coords, axis=2)
    distances_in_AA = distances * lib.param.BOHR

    rho = numpy.array([mopac_param.E2/F03[z] for z in atom_charges])
    gamma = mopac_param.E2 / numpy.sqrt(distances_in_AA**2 +
                                        (rho[:,None] + rho)**2 * .25)
    gamma[numpy.diag_indices(mol.natm)] = 0  # remove self-interaction terms
    return gamma 

def _get_gamma1_half(mol):
    '''gamma1_half[:,i,:] == gamma1[i,:,i,:] == gamma1[i,:,:,i]'''
    atom_charges = mol.atom_charges()
    atom_coords = mol.atom_coords()
    natm = atom_charges.size

    dR = (atom_coords[:,None,:] - atom_coords) * lib.param.BOHR
    distances_in_AA = numpy.linalg.norm(dR, axis=2)

    rho = numpy.array([mopac_param.E2/mopac_param.F03[z] for z in atom_charges])
    div = -mopac_param.E2 / (distances_in_AA**2 + (rho[:,None] + rho)**2*.25)**1.5
    div[numpy.diag_indices(natm)] = 0  # remove self-interaction terms

    gamma1_dense = numpy.einsum('ijx,ij->xij', dR, div)
    return gamma1_dense 

def check_forward(self, diag_data, non_diag_data):
        diag = chainer.Variable(diag_data)
        non_diag = chainer.Variable(non_diag_data)
        y = lower_triangular_matrix(diag, non_diag)

        correct_y = numpy.zeros(
            (self.batch_size, self.n, self.n), dtype=numpy.float32)

        tril_rows, tril_cols = numpy.tril_indices(self.n, -1)
        correct_y[:, tril_rows, tril_cols] = cuda.to_cpu(non_diag_data)

        diag_rows, diag_cols = numpy.diag_indices(self.n)
        correct_y[:, diag_rows, diag_cols] = cuda.to_cpu(diag_data)

        gradient_check.assert_allclose(correct_y, cuda.to_cpu(y.array)) 

def _get_batch_diagonal_cpu(array):
    batch_size, m, n = array.shape
    assert m == n
    rows, cols = np.diag_indices(n)
    return array[:, rows, cols] 

def _set_batch_diagonal_cpu(array, diag_val):
    batch_size, m, n = array.shape
    assert m == n
    rows, cols = np.diag_indices(n)
    array[:, rows, cols] = diag_val 

def _diagonal_idx_array(batch_size, n):
    idx_offsets = np.arange(
        start=0, stop=batch_size * n * n, step=n * n, dtype=np.int32).reshape(
        (batch_size, 1))
    idx = np.ravel_multi_index(
        np.diag_indices(n), (n, n)).reshape((1, n)).astype(np.int32)
    return cuda.to_gpu(idx + idx_offsets) 

def true_backward(A: np.ndarray):
        # TODO: test this with numerical gradient testing!
        I = np.eye(A.shape[1], dtype=floatX)
        idx, idy = np.diag_indices(I)
        return A * (A[..., None] - I[None, ...])[:, idx, idy] 

def _create_lower_upper(self):
        lower = self.lower_entries.new_zeros(self.features, self.features)
        lower[self.lower_indices[0], self.lower_indices[1]] = self.lower_entries
        # The diagonal of L is taken to be all-ones without loss of generality.
        lower[self.diag_indices[0], self.diag_indices[1]] = 1.

        upper = self.upper_entries.new_zeros(self.features, self.features)
        upper[self.upper_indices[0], self.upper_indices[1]] = self.upper_entries
        upper[self.diag_indices[0], self.diag_indices[1]] = self.upper_diag

        return lower, upper 

def __init__(self, features, num_householder, using_cache=False):
        super().__init__(features, using_cache)

        # Parameterization for R
        self.upper_indices = np.triu_indices(features, k=1)
        self.diag_indices = np.diag_indices(features)
        n_triangular_entries = ((features - 1) * features) // 2
        self.upper_entries = nn.Parameter(torch.zeros(n_triangular_entries))
        self.log_upper_diag = nn.Parameter(torch.zeros(features))

        # Parameterization for Q
        self.orthogonal = transforms.HouseholderSequence(
            features=features, num_transforms=num_householder)

        self._initialize()