# Python numpy.diag_indices() Examples

The following are 30 code examples for showing how to use numpy.diag_indices(). These examples are extracted from open source projects. You can vote up the ones you like or vote down the ones you don't like, and go to the original project or source file by following the links above each example.

You may check out the related API usage on the sidebar.

You may also want to check out all available functions/classes of the module , or try the search function .

Example 1
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
scale = alpha = _get_alpha(atom_charges[:,None], atom_charges)

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 
Example 2
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) 
Example 3
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 
Example 4
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) 
Example 5
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 
Example 6
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)

assert np.all(arr[outer_mask] == dst_nodata) 
Example 7
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:
return self.pdf 
Example 8
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 
Example 9
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 
Example 10
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 
Example 11
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]) 
Example 12
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]) 
Example 13
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) 
Example 14
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) 
Example 15
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) 
Example 16
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 
Example 17
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 
Example 18
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 
Example 19
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
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 
Example 20
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
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 
Example 21
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 
Example 22
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 
Example 23
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 
Example 24
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)) 
Example 25
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] 
Example 26
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 
Example 27
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) 
Example 28
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] 
Example 29
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 
Example 30
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()