def invert(self, logvis, einstr='pm,mq'):
'''Return the antenna gain solutions and sky/sep solutions for the provided set of amp/phs measurements.
phs_meas needs to have had any necessary conjugation applied already'''
# Invert the matrices recording which parameters are involved in each measurement
# M_phs is (nmeas,nprm), so M_phs_inv is (nprm,nmeas)
if self.M_phs_inv is None: self.M_phs_inv = n.linalg.pinv(n.array(self.M_phs))
if self.M_amp_inv is None: self.M_amp_inv = n.linalg.pinv(n.array(self.M_amp))
#logvis = n.log(vis)
phs = n.einsum(einstr, self.M_phs_inv, logvis.imag)
amp = n.einsum(einstr, self.M_amp_inv, logvis.real)
g = n.exp(amp + 1j * phs)
i = self.nants * self.npols
g_avg = n.sqrt(n.average(n.abs(g[:i])**2, axis=0))
g[:i] /= n.where(g_avg > 0, g_avg, 1)
g[i:] *= n.abs(g_avg)**2
ant_sol = {}
for ant in self.ants:
ant_sol[ant] = {}
for pol in self.pols:
i = self.antpol_index(ant,pol)
ant_sol[ant][pol] = g[i]
sep_sol = {}
for sep in self.seps:
s = self.sep_index(sep)
sep_sol[sep] = g[s]
return ant_sol, sep_sol
def conv_backward_naive(dout, cache):
"""
A naive implementation of the backward pass for a convolutional layer.
Inputs:
- dout: Upstream derivatives.
- cache: A tuple of (x, w, b, conv_param) as in conv_forward_naive
Returns a tuple of:
- dx: Gradient with respect to x
- dw: Gradient with respect to w
- db: Gradient with respect to b
"""
dx, dw, db = None, None, None
x, w, b, conv_param = cache
stride = conv_param['stride']
pad = conv_param['pad']
N, C, H, W = x.shape
F, _, HH, WW = w.shape
Hp = 1 + (H + 2 * pad - HH) / stride
Wp = 1 + (W + 2 * pad - WW) / stride
dx = np.zeros(x.shape)
dw = np.zeros(w.shape)
db = np.zeros(b.shape)
for i in xrange(N):
# for j in xrange(F):
data = x[i]
data = np.pad(data, ((0, 0), (pad, pad), (pad, pad)), 'constant')
paded_dxi = np.pad(dx[i], ((0, 0), (pad, pad), (pad, pad)), 'constant')
filter_vert_indices = 0
filter_hori_indices = 0
for s in xrange(Hp):
filter_hori_indices = 0
for p in xrange(Wp):
data_fragment = data[:, filter_vert_indices:filter_vert_indices+HH,
filter_hori_indices:filter_hori_indices+WW]
dw += np.einsum('i, jkl->ijkl', dout[i, :, s, p], data_fragment)
# paded_dxi[:, filter_vert_indices:filter_vert_indices+HH,
# filter_hori_indices:filter_hori_indices+WW] = \
# np.einsum('ijkl,i->jkl', w, dout[i, :, s, p])
# paded_dxi[:, filter_vert_indices:filter_vert_indices+HH,
# filter_hori_indices:filter_hori_indices+WW] = \
# np.tensordot(w, dout[i, :, s, p], axes = ([0], [0]))
for f in xrange(F):
paded_dxi[:, filter_vert_indices:filter_vert_indices+HH,
filter_hori_indices:filter_hori_indices+WW] \
+= w[f] * dout[i, f, s, p]
filter_hori_indices += stride
filter_vert_indices += stride
dx[i] = paded_dxi[:, pad:-pad, pad:-pad]
db = np.einsum('ijkl->j', dout)
# print(dx)
#############################################################################
# TODO: Implement the convolutional backward pass. #
#############################################################################
#############################################################################
# END OF YOUR CODE #
#############################################################################
return dx, dw, db
def transform(self, axes, tol=1e-8):
"""
Transforms the elastic constant matrix based on the supplied axes.
Parameters
----------
axes : numpy.ndarray
(3, 3) array giving three right-handed orthogonal vectors to use
for transforming.
tol : float, optional
Relative tolerance to use in identifying near-zero terms.
Returns
-------
ElasticConstants
A new ElasticConstants object that has been transformed.
"""
axes = np.asarray(axes, dtype='float64')
T = axes_check(axes)
Q = np.einsum('km,ln->mnkl', T, T)
C = np.einsum('ghij,ghmn,mnkl->ijkl', Q, self.Cijkl, Q)
C[abs(C / C.max()) < tol] = 0.0
return ElasticConstants(Cijkl=C)
def myprecon(resid, eigval, eigvec):
myprecon_cutoff = 1e-10
local_myprecon = np.zeros([numVars + 1], dtype=float)
for occ in range(numPairs):
for virt in range(numVirt):
denominator = FOCK_mo[numPairs + virt, numPairs + virt] - FOCK_mo[occ, occ] - eigval
if abs(denominator) < myprecon_cutoff:
local_myprecon[occ + numPairs * virt] = eigvec[occ + numPairs * virt] / myprecon_cutoff
else:
# local_myprecon = eigvec / ( diag(H) - eigval ) = K^{-1} u
local_myprecon[occ + numPairs * virt] = eigvec[occ + numPairs * virt] / denominator
if abs(eigval) < myprecon_cutoff:
local_myprecon[numVars] = eigvec[numVars] / myprecon_cutoff
else:
local_myprecon[numVars] = -eigvec[numVars] / eigval
# alpha_myprecon = - ( r, K^{-1} u ) / ( u, K^{-1} u )
alpha_myprecon = -np.einsum("i,i->", local_myprecon, resid) / np.einsum("i,i->", local_myprecon, eigvec)
# local_myprecon = r - ( r, K^{-1} u ) / ( u, K^{-1} u ) * u
local_myprecon = resid + alpha_myprecon * eigvec
for occ in range(numPairs):
for virt in range(numVirt):
denominator = FOCK_mo[numPairs + virt, numPairs + virt] - FOCK_mo[occ, occ] - eigval
if abs(denominator) < myprecon_cutoff:
local_myprecon[occ + numPairs * virt] = -local_myprecon[occ + numPairs * virt] / myprecon_cutoff
else:
local_myprecon[occ + numPairs * virt] = -local_myprecon[occ + numPairs * virt] / denominator
if abs(eigval) < myprecon_cutoff:
local_myprecon[numVars] = -local_myprecon[occ + numPairs * virt] / myprecon_cutoff
else:
local_myprecon[numVars] = local_myprecon[occ + numPairs * virt] / eigval
return local_myprecon
def compute_pixels(orb, sgeom, times, rpy=(0.0, 0.0, 0.0)):
"""Compute cartesian coordinates of the pixels in instrument scan."""
if isinstance(orb, (list, tuple)):
tle1, tle2 = orb
orb = Orbital("mysatellite", line1=tle1, line2=tle2)
# get position and velocity for each time of each pixel
pos, vel = orb.get_position(times, normalize=False)
# now, get the vectors pointing to each pixel
vectors = sgeom.vectors(pos, vel, *rpy)
# compute intersection of lines (directed by vectors and passing through
# (0, 0, 0)) and ellipsoid. Derived from:
# http://en.wikipedia.org/wiki/Line%E2%80%93sphere_intersection
# do the computation between line and ellipsoid (WGS 84)
# NB: AAPP uses GRS 80...
centre = -pos
a__ = 6378.137 # km
# b__ = 6356.75231414 # km, GRS80
b__ = 6356.752314245 # km, WGS84
radius = np.array([[1 / a__, 1 / a__, 1 / b__]]).T
shape = vectors.shape
xr_ = vectors.reshape([3, -1]) * radius
cr_ = centre.reshape([3, -1]) * radius
ldotc = np.einsum("ij,ij->j", xr_, cr_)
lsq = np.einsum("ij,ij->j", xr_, xr_)
csq = np.einsum("ij,ij->j", cr_, cr_)
d1_ = (ldotc - np.sqrt(ldotc ** 2 - csq * lsq + lsq)) / lsq
# return the actual pixel positions
return vectors * d1_.reshape(shape[1:]) - centre
def ITRF_to_GCRS(t, rITRF): # todo: velocity
# Todo: wobble
spin = rot_z(t.gast * tau / 24.0)
position = einsum('ij...,j...->i...', spin, array(rITRF))
return einsum('ij...,j...->i...', t.MT, position)
def analyze(mf, verbose=logger.DEBUG, **kwargs):
'''Analyze the given SCF object: print orbital energies, occupancies;
print orbital coefficients; Mulliken population analysis
'''
from pyscf.lo import orth
from pyscf.tools import dump_mat
mo_energy = mf.mo_energy
mo_occ = mf.mo_occ
mo_coeff = mf.mo_coeff
if isinstance(verbose, logger.Logger):
log = verbose
else:
log = logger.Logger(mf.stdout, verbose)
log.note('**** MO energy ****')
if mf._focka_ao is None:
for i,c in enumerate(mo_occ):
log.note('MO #%-3d energy= %-18.15g occ= %g', i+1, mo_energy[i], c)
else:
mo_ea = numpy.einsum('ik,ik->k', mo_coeff, mf._focka_ao.dot(mo_coeff))
mo_eb = numpy.einsum('ik,ik->k', mo_coeff, mf._fockb_ao.dot(mo_coeff))
log.note(' Roothaan | alpha | beta')
for i,c in enumerate(mo_occ):
log.note('MO #%-3d energy= %-18.15g | %-18.15g | %-18.15g occ= %g',
i+1, mo_energy[i], mo_ea[i], mo_eb[i], c)
ovlp_ao = mf.get_ovlp()
if verbose >= logger.DEBUG:
log.debug(' ** MO coefficients (expansion on meta-Lowdin AOs) **')
label = mf.mol.spheric_labels(True)
orth_coeff = orth.orth_ao(mf.mol, 'meta_lowdin', s=ovlp_ao)
c = reduce(numpy.dot, (orth_coeff.T, ovlp_ao, mo_coeff))
dump_mat.dump_rec(mf.stdout, c, label, start=1, **kwargs)
dm = mf.make_rdm1(mo_coeff, mo_occ)
return mf.mulliken_meta(mf.mol, dm, s=s, verbose=log)
def get_energy(self):
mol, max_iter, ndocc, H, G, A, C, e = self.mol, self.max_iter, self.ndocc,self.H, self.G, self.A, self.C, self.e
E_old = 0.0 # defining initial values
D_old = np.zeros_like(H)
for iteration in range(1, self.max_iter+1): # starting SCF iterations
J = np.einsum('pqrs, rs->pq', G, D_old) # coulomb
K = np.einsum('prqs, rs->pq', G, D_old) # exchange
F = H + J*2 - K # constructing fock matrix
Ft = A.dot(F).dot(A) # transform fock matrix
e, C = np.linalg.eigh(Ft) # diagonalize
C = A.dot(C) # form new SCF eigenvector matrix
Cocc = C[:,:ndocc]
D = np.einsum('pi, qi->pq', Cocc, Cocc) # form new density matrix
E_SCF = np.einsum('pq, pq->', F+H, D) + mol.nuclear_repulsion_energy() # calculate SCF energy
D_norm = np.linalg.norm(D-D_old)
print('RHF iteration {:3d}: energy {:20.14f} dE {:2.5E} D_norm {:2.5E}'.format(iteration, E_SCF, (E_SCF - E_old), D_norm))
if (abs(E_SCF - E_old) < 1.e-10) and D_norm < 1.e-10: # convergence threshold
break
E_old = E_SCF
D_old = D
self.e = e # saving values
self.C = C
self.E_SCF = E_SCF
def get_pp_loc_part1(mydf, cell, kpts):
log = logger.Logger(mydf.stdout, mydf.verbose)
t1 = t0 = (time.clock(), time.time())
nkpts = len(kpts)
gs = mydf.gs
nao = cell.nao_nr()
Gv = cell.get_Gv(gs)
SI = cell.get_SI(Gv)
vpplocG = pseudo.pp_int.get_gth_vlocG_part1(cell, Gv)
vpplocG = -1./cell.vol * numpy.einsum('ij,ij->j', SI, vpplocG)
kpt_allow = numpy.zeros(3)
real = gamma_point(kpts)
if real:
vloc = numpy.zeros((nkpts,nao**2))
else:
vloc = numpy.zeros((nkpts,nao**2), dtype=numpy.complex128)
max_memory = mydf.max_memory - lib.current_memory()[0]
for k, pqkR, pqkI, p0, p1 \
in mydf.ft_loop(cell, mydf.gs, kpt_allow, kpts, max_memory=max_memory):
vG = vpplocG[p0:p1]
if not real:
vloc[k] += numpy.einsum('k,xk->x', vG.real, pqkI) * 1j
vloc[k] += numpy.einsum('k,xk->x', vG.imag, pqkR) *-1j
vloc[k] += numpy.einsum('k,xk->x', vG.real, pqkR)
vloc[k] += numpy.einsum('k,xk->x', vG.imag, pqkI)
pqkR = pqkI = None
t1 = log.timer_debug1('contracting vloc part1', *t1)
return vloc.reshape(-1,nao,nao)
def _assemble(self, mu=None):
g = self.grid
bi = self.boundary_info
if g.dim > 2:
raise NotImplementedError
if bi is None or not bi.has_robin or self.robin_data is None:
return coo_matrix((g.size(g.dim), g.size(g.dim))).tocsc()
RI = bi.robin_boundaries(1)
if g.dim == 1:
robin_c = self.robin_data[0](g.centers(1)[RI], mu=mu)
I = coo_matrix((robin_c, (RI, RI)), shape=(g.size(g.dim), g.size(g.dim)))
return csc_matrix(I).copy()
else:
xref = g.quadrature_points(1, order=self.order)[RI]
# xref(robin-index, quadraturepoint-index)
if self.robin_data[0].shape_range == ():
robin_c = self.robin_data[0](xref, mu=mu)
else:
robin_elements = g.superentities(1, 0)[RI, 0]
robin_indices = g.superentity_indices(1, 0)[RI, 0]
normals = g.unit_outer_normals()[robin_elements, robin_indices]
robin_values = self.robin_data[0](xref, mu=mu)
robin_c = np.einsum('ei,eqi->eq', normals, robin_values)
# robin_c(robin-index, quadraturepoint-index)
q, w = line.quadrature(order=self.order)
SF = np.squeeze(np.array([1 - q, q]))
SF_INTS = np.einsum('ep,pi,pj,e,p->eij', robin_c, SF, SF, g.integration_elements(1)[RI], w).ravel()
SF_I0 = np.repeat(g.subentities(1, g.dim)[RI], 2).ravel()
SF_I1 = np.tile(g.subentities(1, g.dim)[RI], [1, 2]).ravel()
I = coo_matrix((SF_INTS, (SF_I0, SF_I1)), shape=(g.size(g.dim), g.size(g.dim)))
return csc_matrix(I).copy()
def contract_ep(g, fcivec, nsite, nelec, nphonon):
if isinstance(nelec, (int, numpy.number)):
nelecb = nelec//2
neleca = nelec - nelecb
else:
neleca, nelecb = nelec
strsa = numpy.asarray(cistring.gen_strings4orblist(range(nsite), neleca))
strsb = numpy.asarray(cistring.gen_strings4orblist(range(nsite), nelecb))
cishape = make_shape(nsite, nelec, nphonon)
na, nb = cishape[:2]
ci0 = fcivec.reshape(cishape)
fcinew = numpy.zeros(cishape)
nbar = float(neleca+nelecb) / nsite
phonon_cre = numpy.sqrt(numpy.arange(1,nphonon+1))
for i in range(nsite):
maska = (strsa & (1<<i)) > 0
maskb = (strsb & (1<<i)) > 0
e_part = numpy.zeros((na,nb))
e_part[maska,:] += 1
e_part[:,maskb] += 1
e_part[:] -= float(neleca+nelecb) / nsite
for ip in range(nphonon):
slices1 = slices_for_cre(i, nsite, ip)
slices0 = slices_for (i, nsite, ip)
fcinew[slices1] += numpy.einsum('ij...,ij...->ij...', g*phonon_cre[ip]*e_part, ci0[slices0])
fcinew[slices0] += numpy.einsum('ij...,ij...->ij...', g*phonon_cre[ip]*e_part, ci0[slices1])
return fcinew.reshape(fcivec.shape)
def toLab(self, pos, vel=None):
""" Transform the set of coordinates from the magnet frame into the lab
frame.
Parameters:
pos ((n,3) np.ndarray) (mm):
Array of particle positions.
vel ((n,3) np.ndarray) (mm/us):
Array of particle velocities.
"""
#pos = np.atleast_2d(pos)
# Inplace modification
pos[:,0] += self.position[0]
pos[:,1] += self.position[1]
pos[:,2] += self.position[2]
if self.angle != None:
# Generate rotation matrix
rot = np.dot(self._yMatrix(-self.angle[1]),
self._xMatrix(-self.angle[0]))
# Take the dot product of each position and velocity with this matrix.
pos[:] = np.einsum('jk,ik->ij', rot, pos)
if vel != None:
vel[:] = np.einsum('jk,ik->ij', rot, vel)
if vel==None:
return pos
else:
return pos, vel
def hop_uhf2ghf(x1):
x1ab = []
x1ba = []
ip = 0
for k in range(nkpts):
nv = nvira[k]
no = noccb[k]
x1ab.append(x1[ip:ip+nv*no].reshape(nv,no))
ip += nv * no
for k in range(nkpts):
nv = nvirb[k]
no = nocca[k]
x1ba.append(x1[ip:ip+nv*no].reshape(nv,no))
ip += nv * no
dm1ab = []
dm1ba = []
for k in range(nkpts):
d1ab = reduce(numpy.dot, (orbva[k], x1ab[k], orbob[k].T.conj()))
d1ba = reduce(numpy.dot, (orbvb[k], x1ba[k], orboa[k].T.conj()))
dm1ab.append(d1ab+d1ba.T.conj())
dm1ba.append(d1ba+d1ab.T.conj())
v1ao = vresp1(lib.asarray([dm1ab,dm1ba]))
x2ab = [0] * nkpts
x2ba = [0] * nkpts
for k in range(nkpts):
x2ab[k] = numpy.einsum('pr,rq->pq', fvva[k], x1ab[k])
x2ab[k]-= numpy.einsum('sq,ps->pq', foob[k], x1ab[k])
x2ba[k] = numpy.einsum('pr,rq->pq', fvvb[k], x1ba[k])
x2ba[k]-= numpy.einsum('qs,ps->pq', fooa[k], x1ba[k])
x2ab[k] += reduce(numpy.dot, (orbva[k].T.conj(), v1ao[0][k], orbob[k]))
x2ba[k] += reduce(numpy.dot, (orbvb[k].T.conj(), v1ao[1][k], orboa[k]))
return numpy.hstack([x.real.ravel() for x in (x2ab+x2ba)])
def TEME_to_ITRF(jd_ut1, rTEME, vTEME, xp=0.0, yp=0.0):
"""Convert TEME position and velocity into standard ITRS coordinates.
This converts a position and velocity vector in the idiosyncratic
True Equator Mean Equinox (TEME) frame of reference used by the SGP4
theory into vectors into the more standard ITRS frame of reference.
The velocity should be provided in units per day, not per second.
From AIAA 2006-6753 Appendix C.
"""
theta, theta_dot = theta_GMST1982(jd_ut1)
zero = theta_dot * 0.0
angular_velocity = array([zero, zero, -theta_dot])
R = rot_z(-theta)
if len(rTEME.shape) == 1:
rPEF = (R).dot(rTEME)
vPEF = (R).dot(vTEME) + cross(angular_velocity, rPEF)
else:
rPEF = einsum('ij...,j...->i...', R, rTEME)
vPEF = einsum('ij...,j...->i...', R, vTEME) + cross(
angular_velocity, rPEF, 0, 0).T
if xp == 0.0 and yp == 0.0:
rITRF = rPEF
vITRF = vPEF
else:
W = (rot_x(yp)).dot(rot_y(xp))
rITRF = (W).dot(rPEF)
vITRF = (W).dot(vPEF)
return rITRF, vITRF
def toMagnet(self, pos, vel=None):
""" Transform the set of lab-frame coordinates into the coordinate
frame of this array. Positions are shifted then rotated, then velocity
vectors are rotated.
Parameters:
pos ((n,3) np.ndarray) (mm):
Array of particle positions.
vel ((n,3) np.ndarray) (mm/us):
Array of particle velocities.
"""
#pos = np.atleast_2d(pos)
# Move to magnet origin
pos[:,0] -= self.position[0]
pos[:,1] -= self.position[1]
pos[:,2] -= self.position[2]
if self.angle != None:
# Generate rotation matrix
rot = np.dot(self._yMatrix(self.angle[1]), self._xMatrix(self.angle[0]))
# Take the dot product of each position and velocity with this matrix.
pos[:] = np.einsum('jk,ik->ij', rot, pos)
if vel != None:
vel[:] = np.einsum('jk,ik->ij', rot, vel)
if vel==None:
return pos
else:
return pos, vel
def backward(self, dy):
'''
Args:
dy (CTensor): data for the dL / dy, L is the loss
Returns:
dx (Ctensor): data for the dL / dx, L is the loss,
x is the input of current Opertion
'''
# calculations are made on numpy array
if self.axis == 1:
dy = singa.DefaultTranspose(dy)
grad = ctensor2numpy(dy)
output = ctensor2numpy(self.output)
out_1 = np.einsum('ki,ki->ki', grad, output)
medium_out = np.einsum('ki,kj->kij', output, output)
out_2 = np.einsum('kij,kj->ki', medium_out, grad)
out = out_1 - out_2
dx = CTensor(out_1.shape)
dx.CopyFloatDataFromHostPtr(out.flatten())
'''grad = Tensor(data=dy)
output = Tensor(data=self.output)
out_1 = einsum('ki,ki->ki', grad, output)
medium_out = einsum('ki,kj->kij', output, output)
out_2 = einsum('kij,kj->ki', medium_out, grad)
out = out_1 - out_2
dx = CTensor(out_1.data.shape)
dx.CopyFloatDataFromHostPtr(out.data.flatten())'''
if self.axis == 0:
return dx
elif self.axis == 1:
return singa.DefaultTranspose(dx)
def G_Dyson(G0, SigmaDeltaT, Sigma, map):
Beta=map.Beta
G=weight.Weight("SmoothT", map, "TwoSpins", "AntiSymmetric", "K","W")
G0.FFT("K", "W")
SigmaDeltaT.FFT("K")
Sigma.FFT("K", "W")
NSpin, NSub=G.NSpin, G.NSublat
G0SigmaDeltaT=np.einsum("ijklvt,klmnv->ijmnvt",G0.Data, SigmaDeltaT.Data)
G0Sigma=np.einsum("ijklvt,klmnvt->ijmnvt",G0.Data, Sigma.Data)
####correction term
for tau in range(map.MaxTauBin):
G0SigmaDeltaT[...,tau]*= np.cos(np.pi*map.IndexToTau(tau)/Beta)
GS = Beta/map.MaxTauBin*(Beta/map.MaxTauBin*G0Sigma+G0SigmaDeltaT)
#GS shape: NSpin,NSub,NSpin,NSub,Vol,Tau
I=np.eye(NSpin*NSub).reshape([NSpin,NSub,NSpin,NSub])
Denorm=I[...,np.newaxis,np.newaxis]-GS
lu_piv,Determ=weight.LUFactor(Denorm)
Check_Denorminator(Denorm, Determ,map)
G.LUSolve(lu_piv, G0.Data);
return G
def get_compliance_expansion(self):
"""
Gets a compliance tensor expansion from the elastic
tensor expansion.
"""
# TODO: this might have a general form
if not self.order <= 4:
raise ValueError("Compliance tensor expansion only "
"supported for fourth-order and lower")
ce_exp = [ElasticTensor(self[0]).compliance_tensor]
einstring = "ijpq,pqrsuv,rskl,uvmn->ijklmn"
ce_exp.append(np.einsum(einstring, -ce_exp[-1], self[1],
ce_exp[-1], ce_exp[-1]))
if self.order == 4:
# Four terms in the Fourth-Order compliance tensor
einstring_1 = "pqab,cdij,efkl,ghmn,abcdefgh"
tensors_1 = [ce_exp[0]]*4 + [self[-1]]
temp = -np.einsum(einstring_1, *tensors_1)
einstring_2 = "pqab,abcdef,cdijmn,efkl"
einstring_3 = "pqab,abcdef,efklmn,cdij"
einstring_4 = "pqab,abcdef,cdijkl,efmn"
for es in [einstring_2, einstring_3, einstring_4]:
temp -= np.einsum(es, ce_exp[0], self[-2], ce_exp[1], ce_exp[0])
ce_exp.append(temp)
return TensorCollection(ce_exp)
def kernel(cc, eris, t1=None, t2=None, max_memory=2000, verbose=logger.INFO):
assert(isinstance(eris, gccsd._PhysicistsERIs))
if t1 is None or t2 is None:
t1, t2 = cc.t1, cc.t2
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.fock.diagonal().real
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
# e4 = numpy.einsum('ijkabc,ijkabc,ijkabc', t3c.conj(), d3, t3c) / 36
# sia = numpy.einsum('jkbc,ijkabc->ia', eris.oovv, t3c) * .25
# e5 = numpy.einsum('ia,ia', sia, t1.conj())
# et = e4 + e5
# return et
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
et = numpy.einsum('ijkabc,ijkabc,ijkabc', (t3c+t3d).conj(), d3, t3c) / 36
return et
def h_op(x):
x = x.reshape(nvir,nocc)
x2 =-numpy.einsum('sq,ps->pq', foo, x) * 2
x2+= numpy.einsum('pr,rq->pq', fvv, x) * 2
d1 = reduce(numpy.dot, (mo_coeff[:,viridx], x, mo_coeff[:,occidx].T))
if hasattr(mf, 'xc'):
if APPROX_XC_HESSIAN:
vj, vk = mf.get_jk(mol, d1+d1.T)
if abs(hyb) < 1e-10:
dvhf = vj
else:
dvhf = vj - vk * hyb * .5
else:
if save_for_dft[0] is None:
save_for_dft[0] = mf.make_rdm1(mo_coeff, mo_occ)
save_for_dft[1] = mf.get_veff(mol, save_for_dft[0])
dm1 = save_for_dft[0] + d1 + d1.T
vhf1 = mf.get_veff(mol, dm1, dm_last=save_for_dft[0],
vhf_last=save_for_dft[1])
dvhf = vhf1 - save_for_dft[1]
save_for_dft[0] = dm1
save_for_dft[1] = vhf1
else:
dvhf = mf.get_veff(mol, d1+d1.T)
x2 += reduce(numpy.dot, (mo_coeff[:,viridx].T, dvhf,
mo_coeff[:,occidx])) * 4
return x2.reshape(-1)
def scalar(N, Y, fft_form=fft_form_default):
dim = np.size(N)
N = np.array(N, dtype=np.int)
xi = Grid.get_freq(N, Y, fft_form=fft_form)
N_fft=tuple(xi[i].size for i in range(dim))
hGrad = np.zeros((dim,)+N_fft) # zero initialize
for ind in itertools.product(*[range(n) for n in N_fft]):
for i in range(dim):
hGrad[i][ind] = xi[i][ind[i]]
kok= np.einsum('i...,j...->ij...', hGrad, hGrad).real
k2 = np.einsum('i...,i...', hGrad, hGrad).real
ind_center=mean_index(N, fft_form=fft_form)
k2[ind_center]=1.
G0lval=np.zeros_like(kok)
Ival=np.zeros_like(kok)
for ii in range(dim): # diagonal components
G0lval[ii, ii][ind_center] = 1
Ival[ii, ii] = 1
G1l=Tensor(name='G1', val=kok/k2, order=2, N=N, Y=Y, multype=21, Fourier=True, fft_form=fft_form)
G0l=Tensor(name='G1', val=G0lval, order=2, N=N, Y=Y, multype=21, Fourier=True, fft_form=fft_form)
I = Tensor(name='I', val=Ival, order=2, N=N, Y=Y, multype=21, Fourier=True, fft_form=fft_form)
G2l=I-G1l-G0l
return G0l, G1l, G2l
def grad_elec(grad_mf, mo_energy=None, mo_coeff=None, mo_occ=None, atmlst=None):
mf = grad_mf._scf
mol = grad_mf.mol
if mo_energy is None: mo_energy = mf.mo_energy
if mo_occ is None: mo_occ = mf.mo_occ
if mo_coeff is None: mo_coeff = mf.mo_coeff
log = logger.Logger(grad_mf.stdout, grad_mf.verbose)
h1 = grad_mf.get_hcore(mol)
s1 = grad_mf.get_ovlp(mol)
dm0 = mf.make_rdm1(mo_coeff, mo_occ)
t0 = (time.clock(), time.time())
log.debug('Compute Gradients of NR Hartree-Fock Coulomb repulsion')
vhf = grad_mf.get_veff(mol, dm0)
log.timer('gradients of 2e part', *t0)
f1 = h1 + vhf
dme0 = grad_mf.make_rdm1e(mo_energy, mo_coeff, mo_occ)
if atmlst is None:
atmlst = range(mol.natm)
offsetdic = mol.offset_nr_by_atom()
de = numpy.zeros((len(atmlst),3))
for k, ia in enumerate(atmlst):
shl0, shl1, p0, p1 = offsetdic[ia]
# h1, s1, vhf are \nabla <i|h|j>, the nuclear gradients = -\nabla
vrinv = grad_mf._grad_rinv(mol, ia)
de[k] += numpy.einsum('xij,ij->x', f1[:,p0:p1], dm0[p0:p1]) * 2
de[k] += numpy.einsum('xij,ij->x', vrinv, dm0) * 2
de[k] -= numpy.einsum('xij,ij->x', s1[:,p0:p1], dme0[p0:p1]) * 2
log.debug('gradients of electronic part')
log.debug(str(de))
return de
def compute_inertia_tensor(traj):
"""Compute the inertia tensor of a trajectory.
For each frame,
I_{ab} = sum_{i_atoms} [m_i * (r_i^2 * d_{ab} - r_{ia} * r_{ib})]
Parameters
----------
traj : Trajectory
Trajectory to compute inertia tensor of.
Returns
-------
I_ab: np.ndarray, shape=(traj.n_frames, 3, 3), dtype=float64
Inertia tensors for each frame.
"""
center_of_mass = np.expand_dims(compute_center_of_mass(traj), axis=1)
xyz = traj.xyz - center_of_mass
masses = np.array([atom.element.mass for atom in traj.top.atoms])
eyes = np.empty(shape=(traj.n_frames, 3, 3), dtype=np.float64)
eyes[:] = np.eye(3)
A = np.einsum("i, kij->k", masses, xyz ** 2).reshape(traj.n_frames, 1, 1)
B = np.einsum("ij..., ...jk->...ki", masses[:, np.newaxis] * xyz.T, xyz)
return A * eyes - B
def Srsi(mc, dms, eris, verbose=None):
#Subspace S_ijr^{(1)}
mo_core, mo_cas, mo_virt = _extract_orbs(mc, mc.mo_coeff)
dm1 = dms['1']
dm2 = dms['2']
ncore = mo_core.shape[1]
ncas = mo_cas.shape[1]
nocc = ncore + ncas
if eris is None:
h1e = mc.h1e_for_cas()[0]
h2e = ao2mo.restore(1, mc.ao2mo(mo_cas), ncas).transpose(0,2,1,3)
h2e_v = ao2mo.incore.general(mc._scf._eri,[mo_virt,mo_core,mo_virt,mo_cas],compact=False)
h2e_v = h2e_v.reshape(mo_virt.shape[1],ncore,mo_virt.shape[1],ncas).transpose(0,2,1,3)
else:
h1e = eris['h1eff'][ncore:nocc,ncore:nocc]
h2e = eris['ppaa'][ncore:nocc,ncore:nocc].transpose(0,2,1,3)
h2e_v = eris['pacv'][nocc:].transpose(3,0,2,1)
k27 = make_k27(h1e,h2e,dm1,dm2)
norm = 2.0*numpy.einsum('rsip,rsia,pa->rsi',h2e_v,h2e_v,dm1)\
- 1.0*numpy.einsum('rsip,sria,pa->rsi',h2e_v,h2e_v,dm1)
h = 2.0*numpy.einsum('rsip,rsia,pa->rsi',h2e_v,h2e_v,k27)\
- 1.0*numpy.einsum('rsip,sria,pa->rsi',h2e_v,h2e_v,k27)
diff = mc.mo_energy[nocc:,None,None] + mc.mo_energy[None,nocc:,None] - mc.mo_energy[None,None,:ncore]
return _norm_to_energy(norm, h, diff)
def Srs(mc, dms, eris=None, verbose=None):
#Subspace S_rs^{(-2)}
mo_core, mo_cas, mo_virt = _extract_orbs(mc, mc.mo_coeff)
dm1 = dms['1']
dm2 = dms['2']
dm3 = dms['3']
ncore = mo_core.shape[1]
ncas = mo_cas.shape[1]
nocc = ncore + ncas
if mo_virt.shape[1] ==0:
return 0, 0
if eris is None:
h1e = mc.h1e_for_cas()[0]
h2e = ao2mo.restore(1, mc.ao2mo(mo_cas), ncas).transpose(0,2,1,3)
h2e_v = ao2mo.incore.general(mc._scf._eri,[mo_virt,mo_cas,mo_virt,mo_cas],compact=False)
h2e_v = h2e_v.reshape(mo_virt.shape[1],ncas,mo_virt.shape[1],ncas).transpose(0,2,1,3)
else:
h1e = eris['h1eff'][ncore:nocc,ncore:nocc]
h2e = eris['ppaa'][ncore:nocc,ncore:nocc].transpose(0,2,1,3)
h2e_v = eris['papa'][nocc:,:,nocc:].transpose(0,2,1,3)
# a7 is very sensitive to the accuracy of HF orbital and CI wfn
rm2, a7 = make_a7(h1e,h2e,dm1,dm2,dm3)
norm = 0.5*numpy.einsum('rsqp,rsba,pqba->rs',h2e_v,h2e_v,rm2)
h = 0.5*numpy.einsum('rsqp,rsba,pqab->rs',h2e_v,h2e_v,a7)
diff = mc.mo_energy[nocc:,None] + mc.mo_energy[None,nocc:]
return _norm_to_energy(norm, h, diff)
def Sijrs(mc, eris, verbose=None):
mo_core, mo_cas, mo_virt = _extract_orbs(mc, mc.mo_coeff)
ncore = mo_core.shape[1]
nvirt = mo_virt.shape[1]
ncas = mo_cas.shape[1]
nocc = ncore + ncas
if eris is None:
erifile = tempfile.NamedTemporaryFile(dir=lib.param.TMPDIR)
feri = ao2mo.outcore.general(mc.mol, (mo_core,mo_virt,mo_core,mo_virt),
erifile.name, verbose=mc.verbose)
else:
feri = eris['cvcv']
eia = mc.mo_energy[:ncore,None] -mc.mo_energy[None,nocc:]
norm = 0
e = 0
with ao2mo.load(feri) as cvcv:
for i in range(ncore):
djba = (eia.reshape(-1,1) + eia[i].reshape(1,-1)).ravel()
gi = numpy.asarray(cvcv[i*nvirt:(i+1)*nvirt])
gi = gi.reshape(nvirt,ncore,nvirt).transpose(1,2,0)
t2i = (gi.ravel()/djba).reshape(ncore,nvirt,nvirt)
# 2*ijab-ijba
theta = gi*2 - gi.transpose(0,2,1)
norm += numpy.einsum('jab,jab', gi, theta)
e += numpy.einsum('jab,jab', t2i, theta)
return norm, e
def Sijr(mc, dms, eris, verbose=None):
#Subspace S_ijr^{(1)}
mo_core, mo_cas, mo_virt = _extract_orbs(mc, mc.mo_coeff)
dm1 = dms['1']
dm2 = dms['2']
ncore = mo_core.shape[1]
ncas = mo_cas.shape[1]
nocc = ncore + ncas
if eris is None:
h1e = mc.h1e_for_cas()[0]
h2e = ao2mo.restore(1, mc.ao2mo(mo_cas), ncas).transpose(0,2,1,3)
h2e_v = ao2mo.incore.general(mc._scf._eri,[mo_virt,mo_core,mo_cas,mo_core],compact=False)
h2e_v = h2e_v.reshape(mo_virt.shape[1],ncore,ncas,ncore).transpose(0,2,1,3)
else:
h1e = eris['h1eff'][ncore:nocc,ncore:nocc]
h2e = eris['ppaa'][ncore:nocc,ncore:nocc].transpose(0,2,1,3)
h2e_v = eris['pacv'][:ncore].transpose(3,1,2,0)
if 'h1' in dms:
hdm1 = dms['h1']
else:
hdm1 = make_hdm1(dm1)
a3 = make_a3(h1e,h2e,dm1,dm2,hdm1)
norm = 2.0*numpy.einsum('rpji,raji,pa->rji',h2e_v,h2e_v,hdm1)\
- 1.0*numpy.einsum('rpji,raij,pa->rji',h2e_v,h2e_v,hdm1)
h = 2.0*numpy.einsum('rpji,raji,pa->rji',h2e_v,h2e_v,a3)\
- 1.0*numpy.einsum('rpji,raij,pa->rji',h2e_v,h2e_v,a3)
diff = mc.mo_energy[nocc:,None,None] - mc.mo_energy[None,:ncore,None] - mc.mo_energy[None,None,:ncore]
return _norm_to_energy(norm, h, diff)
def einsum(subscripts, *tensors, **kwargs):
'''Perform a more efficient einsum via reshaping to a matrix multiply.
Current differences compared to numpy.einsum:
This assumes that each repeated index is actually summed (i.e. no 'i,i->i')
and appears only twice (i.e. no 'ij,ik,il->jkl'). The output indices must
be explicitly specified (i.e. 'ij,j->i' and not 'ij,j').
'''
contract = kwargs.pop('_contract', _contract)
subscripts = subscripts.replace(' ','')
if len(tensors) <= 1 or '...' in subscripts:
out = numpy.einsum(subscripts, *tensors, **kwargs)
elif len(tensors) <= 2:
out = _contract(subscripts, *tensors, **kwargs)
else:
if '->' in subscripts:
indices_in, idx_final = subscripts.split('->')
indices_in = indices_in.split(',')
else:
idx_final = ''
indices_in = subscripts.split('->')[0].split(',')
tensors = list(tensors)
contraction_list = _einsum_path(subscripts, *tensors, optimize=True,
einsum_call=True)[1]
for contraction in contraction_list:
inds, idx_rm, einsum_str, remaining = contraction[:4]
tmp_operands = [tensors.pop(x) for x in inds]
if len(tmp_operands) > 2:
out = numpy.einsum(einsum_str, *tmp_operands)
else:
out = contract(einsum_str, *tmp_operands)
tensors.append(out)
return out
def make_a23(h1e,h2e,dm1,dm2,dm3):
a23 = -numpy.einsum('ip,caib->abcp',h1e,dm2)\
-numpy.einsum('pijk,cajbik->abcp',h2e,dm3)\
+2.0*numpy.einsum('bp,ca->abcp',h1e,dm1)\
+2.0*numpy.einsum('pibk,caik->abcp',h2e,dm2)
return a23
def h_op(x):
x1 = numpy.zeros_like(focka)
x1[mask] = x
x1 = x1 - x1.T
x2 = numpy.zeros_like(focka)
#: x2[nb:,:na] = numpy.einsum('sp,qs->pq', focka[:na,nb:], x1[:na,:na])
#: x2[nb:,:na] += numpy.einsum('rq,rp->pq', focka[:na,:na], x1[:na,nb:])
#: x2[na:,:na] -= numpy.einsum('sp,rp->rs', focka[:na,:na], x1[na:,:na])
#: x2[na:,:na] -= numpy.einsum('ps,qs->pq', focka[na:], x1[:na]) * 2
#: x2[nb:na,:nb] += numpy.einsum('qr,pr->pq', focka[:nb], x1[nb:na])
#: x2[nb:na,:nb] -= numpy.einsum('rq,sq->rs', focka[nb:na], x1[:nb])
#: x2[nb:,:na] += numpy.einsum('sp,qs->pq', fockb[:nb,nb:], x1[:na,:nb])
#: x2[nb:,:na] += numpy.einsum('rq,rp->pq', fockb[:nb,:na], x1[:nb,nb:])
#: x2[nb:,:nb] -= numpy.einsum('sp,rp->rs', fockb[:nb], x1[nb:])
#: x2[nb:,:nb] -= numpy.einsum('rq,sq->rs', fockb[nb:], x1[:nb]) * 2
x2[viridxb[:,None],occidxa] = \
(numpy.einsum('sp,qs->pq', focka[occidxa[:,None],viridxb], x1[occidxa[:,None],occidxa])
+numpy.einsum('rq,rp->pq', focka[occidxa[:,None],occidxa], x1[occidxa[:,None],viridxb]))
x2[viridxa[:,None],occidxa] -= \
(numpy.einsum('sp,rp->rs', focka[occidxa[:,None],occidxa], x1[viridxa[:,None],occidxa])
+numpy.einsum('ps,qs->pq', focka[viridxa], x1[occidxa]) * 2)
x2[occidxa[:,None],occidxb] += \
(numpy.einsum('qr,pr->pq', focka[occidxb], x1[occidxa])
-numpy.einsum('rq,sq->rs', focka[occidxa], x1[occidxb]))
x2[viridxb[:,None],occidxa] += \
(numpy.einsum('sp,qs->pq', fockb[occidxb[:,None],viridxb], x1[occidxa[:,None],occidxb])
+numpy.einsum('rq,rp->pq', fockb[occidxb[:,None],occidxa], x1[occidxb[:,None],viridxb]))
x2[viridxb[:,None],occidxb] -= \
(numpy.einsum('sp,rp->rs', fockb[occidxb], x1[viridxb])
+numpy.einsum('rq,sq->rs', fockb[viridxb], x1[occidxb]) * 2)
x2 *= .5
d1a = reduce(numpy.dot, (mo_coeff[:,viridxa], x1[viridxa[:,None],occidxa], mo_coeff[:,occidxa].T))
d1b = reduce(numpy.dot, (mo_coeff[:,viridxb], x1[viridxb[:,None],occidxb], mo_coeff[:,occidxb].T))
if hasattr(mf, 'xc'):
if APPROX_XC_HESSIAN:
vj, vk = mf.get_jk(mol, numpy.array((d1a+d1a.T,d1b+d1b.T)))
if abs(hyb) < 1e-10:
dvhf = vj[0] + vj[1]
else:
dvhf = (vj[0] + vj[1]) - vk * hyb
else:
from pyscf.dft import uks
if save_for_dft[0] is None:
save_for_dft[0] = mf.make_rdm1(mo_coeff, mo_occ)
save_for_dft[1] = mf.get_veff(mol, save_for_dft[0])
dm1 = numpy.array((save_for_dft[0][0]+d1a+d1a.T,
save_for_dft[0][1]+d1b+d1b.T))
vhf1 = uks.get_veff_(mf, mol, dm1, dm_last=save_for_dft[0],
vhf_last=save_for_dft[1])
dvhf = (vhf1[0]-save_for_dft[1][0], vhf1[1]-save_for_dft[1][1])
save_for_dft[0] = dm1
save_for_dft[1] = vhf1
else:
dvhf = mf.get_veff(mol, numpy.array((d1a+d1a.T,d1b+d1b.T)))
x2[viridxa[:,None],occidxa] += reduce(numpy.dot, (mo_coeff[:,viridxa].T, dvhf[0], mo_coeff[:,occidxa]))
x2[viridxb[:,None],occidxb] += reduce(numpy.dot, (mo_coeff[:,viridxb].T, dvhf[1], mo_coeff[:,occidxb]))
return x2[mask]
def energy(cc, t1=None, t2=None, eris=None):
'''UCCSD correlation energy'''
if t1 is None: t1 = cc.t1
if t2 is None: t2 = cc.t2
if eris is None: eris = cc.ao2mo()
t1a, t1b = t1
t2aa, t2ab, t2bb = t2
nocca, noccb, nvira, nvirb = t2ab.shape
eris_ovov = np.asarray(eris.ovov)
eris_OVOV = np.asarray(eris.OVOV)
eris_ovOV = np.asarray(eris.ovOV)
fova = eris.focka[:nocca, nocca:]
fovb = eris.fockb[:noccb, noccb:]
e = np.einsum('ia,ia', fova, t1a)
e += np.einsum('ia,ia', fovb, t1b)
e += 0.25 * np.einsum('ijab,iajb', t2aa, eris_ovov)
e -= 0.25 * np.einsum('ijab,ibja', t2aa, eris_ovov)
e += 0.25 * np.einsum('ijab,iajb', t2bb, eris_OVOV)
e -= 0.25 * np.einsum('ijab,ibja', t2bb, eris_OVOV)
e += np.einsum('iJaB,iaJB', t2ab, eris_ovOV)
e += 0.5 * np.einsum('ia,jb,iajb', t1a, t1a, eris_ovov)
e -= 0.5 * np.einsum('ia,jb,ibja', t1a, t1a, eris_ovov)
e += 0.5 * np.einsum('ia,jb,iajb', t1b, t1b, eris_OVOV)
e -= 0.5 * np.einsum('ia,jb,ibja', t1b, t1b, eris_OVOV)
e += np.einsum('ia,jb,iajb', t1a, t1b, eris_ovOV)
if abs(e.imag) > 1e-4:
logger.warn(cc, 'Non-zero imaginary part found in UCCSD energy %s', e)
return e.real
def gen_g_hop(casscf, mo, ci0, eris, verbose=None):
ncas = casscf.ncas
ncore = casscf.ncore
nocc = ncas + ncore
nelecas = casscf.nelecas
nmo = mo.shape[1]
ci0 = ci0.ravel()
if hasattr(casscf.fcisolver, 'gen_linkstr'):
linkstrl = casscf.fcisolver.gen_linkstr(ncas, nelecas, True)
linkstr = casscf.fcisolver.gen_linkstr(ncas, nelecas, False)
else:
linkstrl = linkstr = None
def fci_matvec(civec, h1, h2):
h2cas = casscf.fcisolver.absorb_h1e(h1, h2, ncas, nelecas, .5)
hc = casscf.fcisolver.contract_2e(h2cas,
civec,
ncas,
nelecas,
link_index=linkstrl).ravel()
return hc
# part5
jkcaa = numpy.empty((nocc, ncas))
# part2, part3
vhf_a = numpy.empty((nmo, nmo))
# part1 ~ (J + 2K)
casdm1, casdm2 = casscf.fcisolver.make_rdm12(ci0,
ncas,
nelecas,
link_index=linkstr)
dm2tmp = casdm2.transpose(1, 2, 0, 3) + casdm2.transpose(0, 2, 1, 3)
dm2tmp = dm2tmp.reshape(ncas**2, -1)
hdm2 = numpy.empty((nmo, ncas, nmo, ncas))
g_dm2 = numpy.empty((nmo, ncas))
eri_cas = numpy.empty((ncas, ncas, ncas, ncas))
for i in range(nmo):
jbuf = eris.ppaa[i]
kbuf = eris.papa[i]
if i < nocc:
jkcaa[i] = numpy.einsum('ik,ik->i', 6 * kbuf[:, i] - 2 * jbuf[i],
casdm1)
vhf_a[i] = (numpy.einsum('quv,uv->q', jbuf, casdm1) -
numpy.einsum('uqv,uv->q', kbuf, casdm1) * .5)
jtmp = lib.dot(jbuf.reshape(nmo, -1), casdm2.reshape(ncas * ncas, -1))
jtmp = jtmp.reshape(nmo, ncas, ncas)
ktmp = lib.dot(kbuf.transpose(1, 0, 2).reshape(nmo, -1), dm2tmp)
hdm2[i] = (ktmp.reshape(nmo, ncas, ncas) + jtmp).transpose(1, 0, 2)
g_dm2[i] = numpy.einsum('uuv->v', jtmp[ncore:nocc])
if ncore <= i < nocc:
eri_cas[i - ncore] = jbuf[ncore:nocc]
jbuf = kbuf = jtmp = ktmp = dm2tmp = casdm2 = None
vhf_ca = eris.vhf_c + vhf_a
h1e_mo = reduce(numpy.dot, (mo.T, casscf.get_hcore(), mo))
################# gradient #################
gpq = numpy.zeros_like(h1e_mo)
gpq[:, :ncore] = (h1e_mo[:, :ncore] + vhf_ca[:, :ncore]) * 2
gpq[:, ncore:nocc] = numpy.dot(
h1e_mo[:, ncore:nocc] + eris.vhf_c[:, ncore:nocc], casdm1)
gpq[:, ncore:nocc] += g_dm2
h1cas_0 = h1e_mo[ncore:nocc, ncore:nocc] + eris.vhf_c[ncore:nocc,
ncore:nocc]
h2cas_0 = casscf.fcisolver.absorb_h1e(h1cas_0, eri_cas, ncas, nelecas, .5)
hc0 = casscf.fcisolver.contract_2e(h2cas_0,
ci0,
ncas,
nelecas,
link_index=linkstrl).ravel()
eci0 = ci0.dot(hc0)
gci = hc0 - ci0 * eci0
def g_update(u, fcivec):
uc = u[:, :ncore].copy()
ua = u[:, ncore:nocc].copy()
rmat = u - numpy.eye(nmo)
ra = rmat[:, ncore:nocc].copy()
mo1 = numpy.dot(mo, u)
mo_c = numpy.dot(mo, uc)
mo_a = numpy.dot(mo, ua)
dm_c = numpy.dot(mo_c, mo_c.T) * 2
fcivec *= 1. / numpy.linalg.norm(fcivec)
casdm1, casdm2 = casscf.fcisolver.make_rdm12(fcivec,
ncas,
nelecas,
link_index=linkstr)
#casscf.with_dep4 = False
#casscf.ci_response_space = 3
#casscf.ci_grad_trust_region = 3
#casdm1, casdm2, gci, fcivec = casscf.update_casdm(mo, u, fcivec, 0, eris, locals())
dm_a = reduce(numpy.dot, (mo_a, casdm1, mo_a.T))
vj, vk = casscf.get_jk(casscf.mol, (dm_c, dm_a))
vhf_c = reduce(numpy.dot, (mo1.T, vj[0] - vk[0] * .5, mo1[:, :nocc]))
vhf_a = reduce(numpy.dot, (mo1.T, vj[1] - vk[1] * .5, mo1[:, :nocc]))
h1e_mo1 = reduce(numpy.dot, (u.T, h1e_mo, u[:, :nocc]))
p1aa = numpy.empty((nmo, ncas, ncas * ncas))
paa1 = numpy.empty((nmo, ncas * ncas, ncas))
aaaa = numpy.empty([ncas] * 4)
for i in range(nmo):
jbuf = eris.ppaa[i]
kbuf = eris.papa[i]
p1aa[i] = lib.dot(ua.T, jbuf.reshape(nmo, -1))
paa1[i] = lib.dot(kbuf.transpose(0, 2, 1).reshape(-1, nmo), ra)
if ncore <= i < nocc:
aaaa[i - ncore] = jbuf[ncore:nocc]
# active space Hamiltonian up to 2nd order
aa11 = lib.dot(ua.T, p1aa.reshape(nmo, -1)).reshape([ncas] * 4)
aa11 = aa11 + aa11.transpose(2, 3, 0, 1) - aaaa
a11a = lib.dot(ra.T, paa1.reshape(nmo, -1)).reshape((ncas, ) * 4)
a11a = a11a + a11a.transpose(1, 0, 2, 3)
a11a = a11a + a11a.transpose(0, 1, 3, 2)
eri_cas_2 = aa11 + a11a
h1cas_2 = h1e_mo1[ncore:nocc, ncore:nocc] + vhf_c[ncore:nocc,
ncore:nocc]
fcivec = fcivec.ravel()
hc0 = fci_matvec(fcivec, h1cas_2, eri_cas_2)
gci = hc0 - fcivec * fcivec.dot(hc0)
g = numpy.zeros_like(h1e_mo)
g[:, :ncore] = (h1e_mo1[:, :ncore] + vhf_c[:, :ncore] +
vhf_a[:, :ncore]) * 2
g[:, ncore:nocc] = numpy.dot(
h1e_mo1[:, ncore:nocc] + vhf_c[:, ncore:nocc], casdm1)
# 0000 + 1000 + 0100 + 0010 + 0001 + 1100 + 1010 + 1001 (missing 0110 + 0101 + 0011)
p1aa = lib.dot(u.T, p1aa.reshape(nmo,
-1)).reshape(nmo, ncas, ncas, ncas)
paa1 = lib.dot(u.T, paa1.reshape(nmo,
-1)).reshape(nmo, ncas, ncas, ncas)
p1aa += paa1
p1aa += paa1.transpose(0, 1, 3, 2)
g[:, ncore:nocc] += numpy.einsum('puwx,wxuv->pv', p1aa, casdm2)
g_orb = casscf.pack_uniq_var(g - g.T)
return numpy.hstack((g_orb * 2, gci * 2))
############## hessian, diagonal ###########
# part7
dm1 = numpy.zeros((nmo, nmo))
idx = numpy.arange(ncore)
dm1[idx, idx] = 2
dm1[ncore:nocc, ncore:nocc] = casdm1
h_diag = numpy.einsum('ii,jj->ij', h1e_mo, dm1) - h1e_mo * dm1
h_diag = h_diag + h_diag.T
# part8
g_diag = gpq.diagonal()
h_diag -= g_diag + g_diag.reshape(-1, 1)
idx = numpy.arange(nmo)
h_diag[idx, idx] += g_diag * 2
# part2, part3
v_diag = vhf_ca.diagonal() # (pr|kl) * E(sq,lk)
h_diag[:, :ncore] += v_diag.reshape(-1, 1) * 2
h_diag[:ncore] += v_diag * 2
idx = numpy.arange(ncore)
h_diag[idx, idx] -= v_diag[:ncore] * 4
# V_{pr} E_{sq}
tmp = numpy.einsum('ii,jj->ij', eris.vhf_c, casdm1)
h_diag[:, ncore:nocc] += tmp
h_diag[ncore:nocc, :] += tmp.T
tmp = -eris.vhf_c[ncore:nocc, ncore:nocc] * casdm1
h_diag[ncore:nocc, ncore:nocc] += tmp + tmp.T
# part4
# -2(pr|sq) + 4(pq|sr) + 4(pq|rs) - 2(ps|rq)
tmp = 6 * eris.k_pc - 2 * eris.j_pc
h_diag[ncore:, :ncore] += tmp[ncore:]
h_diag[:ncore, ncore:] += tmp[ncore:].T
# part5 and part6 diag
# -(qr|kp) E_s^k p in core, sk in active
h_diag[:nocc, ncore:nocc] -= jkcaa
h_diag[ncore:nocc, :nocc] -= jkcaa.T
v_diag = numpy.einsum('ijij->ij', hdm2)
h_diag[ncore:nocc, :] += v_diag.T
h_diag[:, ncore:nocc] += v_diag
# Does this term contribute to internal rotation?
# h_diag[ncore:nocc,ncore:nocc] -= v_diag[:,ncore:nocc]*2
h_diag = casscf.pack_uniq_var(h_diag)
hci_diag = casscf.fcisolver.make_hdiag(h1cas_0, eri_cas, ncas, nelecas)
hci_diag -= eci0
hci_diag -= gci * ci0 * 4
hdiag_all = numpy.hstack((h_diag * 2, hci_diag * 2))
g_orb = casscf.pack_uniq_var(gpq - gpq.T)
g_all = numpy.hstack((g_orb * 2, gci * 2))
ngorb = g_orb.size
def h_op(x):
x1 = casscf.unpack_uniq_var(x[:ngorb])
ci1 = x[ngorb:]
# H_cc
hci1 = casscf.fcisolver.contract_2e(h2cas_0,
ci1,
ncas,
nelecas,
link_index=linkstrl).ravel()
hci1 -= ci1 * eci0
hci1 -= ((hc0 - ci0 * eci0) * ci0.dot(ci1) + ci0 *
(hc0 - ci0 * eci0).dot(ci1)) * 2
# H_co
rc = x1[:, :ncore]
ra = x1[:, ncore:nocc]
ddm_c = numpy.zeros((nmo, nmo))
ddm_c[:, :ncore] = rc[:, :ncore] * 2
ddm_c[:ncore, :] += rc[:, :ncore].T * 2
tdm1, tdm2 = casscf.fcisolver.trans_rdm12(ci1,
ci0,
ncas,
nelecas,
link_index=linkstr)
tdm1 = tdm1 + tdm1.T
tdm2 = tdm2 + tdm2.transpose(1, 0, 3, 2)
tdm2 = (tdm2 + tdm2.transpose(2, 3, 0, 1)) * .5
vhf_a = numpy.empty((nmo, ncore))
paaa = numpy.empty((nmo, ncas, ncas, ncas))
jk = 0
for i in range(nmo):
jbuf = eris.ppaa[i]
kbuf = eris.papa[i]
paaa[i] = jbuf[ncore:nocc]
vhf_a[i] = numpy.einsum('quv,uv->q', jbuf[:ncore], tdm1)
vhf_a[i] -= numpy.einsum('uqv,uv->q', kbuf[:, :ncore], tdm1) * .5
jk += numpy.einsum('quv,q->uv', jbuf, ddm_c[i])
jk -= numpy.einsum('uqv,q->uv', kbuf, ddm_c[i]) * .5
g_dm2 = numpy.einsum('puwx,wxuv->pv', paaa, tdm2)
aaaa = numpy.dot(ra.T, paaa.reshape(nmo, -1)).reshape([ncas] * 4)
aaaa = aaaa + aaaa.transpose(1, 0, 2, 3)
aaaa = aaaa + aaaa.transpose(2, 3, 0, 1)
h1aa = numpy.dot(h1e_mo[ncore:nocc] + eris.vhf_c[ncore:nocc], ra)
h1aa = h1aa + h1aa.T + jk
h1c0 = fci_matvec(ci0, h1aa, aaaa)
hci1 += h1c0
hci1 -= h1c0.dot(ci0) * ci0
# H_oo
# part7
# (-h_{sp} R_{rs} gamma_{rq} - h_{rq} R_{pq} gamma_{sp})/2 + (pr<->qs)
x2 = reduce(lib.dot, (h1e_mo, x1, dm1))
# part8
# (g_{ps}\delta_{qr}R_rs + g_{qr}\delta_{ps}) * R_pq)/2 + (pr<->qs)
x2 -= numpy.dot((gpq + gpq.T), x1) * .5
# part2
# (-2Vhf_{sp}\delta_{qr}R_pq - 2Vhf_{qr}\delta_{sp}R_rs)/2 + (pr<->qs)
x2[:ncore] += reduce(numpy.dot,
(x1[:ncore, ncore:], vhf_ca[ncore:])) * 2
# part3
# (-Vhf_{sp}gamma_{qr}R_{pq} - Vhf_{qr}gamma_{sp}R_{rs})/2 + (pr<->qs)
x2[ncore:nocc] += reduce(numpy.dot,
(casdm1, x1[ncore:nocc], eris.vhf_c))
# part1
x2[:, ncore:nocc] += numpy.einsum('purv,rv->pu', hdm2, x1[:,
ncore:nocc])
if ncore > 0:
# part4, part5, part6
# Due to x1_rs [4(pq|sr) + 4(pq|rs) - 2(pr|sq) - 2(ps|rq)] for r>s p>q,
# == -x1_sr [4(pq|sr) + 4(pq|rs) - 2(pr|sq) - 2(ps|rq)] for r>s p>q,
# x2[:,:ncore] += H * x1[:,:ncore] => (becuase x1=-x1.T) =>
# x2[:,:ncore] += -H' * x1[:ncore] => (becuase x2-x2.T) =>
# x2[:ncore] += H' * x1[:ncore]
va, vc = casscf.update_jk_in_ah(mo, x1, casdm1, eris)
x2[ncore:nocc] += va
x2[:ncore, ncore:] += vc
# H_oc
s10 = ci1.dot(ci0) * 2
x2[:, :ncore] += (
(h1e_mo[:, :ncore] + eris.vhf_c[:, :ncore]) * s10 + vhf_a) * 2
x2[:, ncore:nocc] += numpy.dot(
h1e_mo[:, ncore:nocc] + eris.vhf_c[:, ncore:nocc], tdm1)
x2[:, ncore:nocc] += g_dm2
x2 -= s10 * gpq
# (pr<->qs)
x2 = x2 - x2.T
return numpy.hstack((casscf.pack_uniq_var(x2) * 2, hci1 * 2))
return g_all, g_update, h_op, hdiag_all
def h_op(x):
x1 = casscf.unpack_uniq_var(x[:ngorb])
ci1 = x[ngorb:]
# H_cc
hci1 = casscf.fcisolver.contract_2e(h2cas_0,
ci1,
ncas,
nelecas,
link_index=linkstrl).ravel()
hci1 -= ci1 * eci0
hci1 -= ((hc0 - ci0 * eci0) * ci0.dot(ci1) + ci0 *
(hc0 - ci0 * eci0).dot(ci1)) * 2
# H_co
rc = x1[:, :ncore]
ra = x1[:, ncore:nocc]
ddm_c = numpy.zeros((nmo, nmo))
ddm_c[:, :ncore] = rc[:, :ncore] * 2
ddm_c[:ncore, :] += rc[:, :ncore].T * 2
tdm1, tdm2 = casscf.fcisolver.trans_rdm12(ci1,
ci0,
ncas,
nelecas,
link_index=linkstr)
tdm1 = tdm1 + tdm1.T
tdm2 = tdm2 + tdm2.transpose(1, 0, 3, 2)
tdm2 = (tdm2 + tdm2.transpose(2, 3, 0, 1)) * .5
vhf_a = numpy.empty((nmo, ncore))
paaa = numpy.empty((nmo, ncas, ncas, ncas))
jk = 0
for i in range(nmo):
jbuf = eris.ppaa[i]
kbuf = eris.papa[i]
paaa[i] = jbuf[ncore:nocc]
vhf_a[i] = numpy.einsum('quv,uv->q', jbuf[:ncore], tdm1)
vhf_a[i] -= numpy.einsum('uqv,uv->q', kbuf[:, :ncore], tdm1) * .5
jk += numpy.einsum('quv,q->uv', jbuf, ddm_c[i])
jk -= numpy.einsum('uqv,q->uv', kbuf, ddm_c[i]) * .5
g_dm2 = numpy.einsum('puwx,wxuv->pv', paaa, tdm2)
aaaa = numpy.dot(ra.T, paaa.reshape(nmo, -1)).reshape([ncas] * 4)
aaaa = aaaa + aaaa.transpose(1, 0, 2, 3)
aaaa = aaaa + aaaa.transpose(2, 3, 0, 1)
h1aa = numpy.dot(h1e_mo[ncore:nocc] + eris.vhf_c[ncore:nocc], ra)
h1aa = h1aa + h1aa.T + jk
h1c0 = fci_matvec(ci0, h1aa, aaaa)
hci1 += h1c0
hci1 -= h1c0.dot(ci0) * ci0
# H_oo
# part7
# (-h_{sp} R_{rs} gamma_{rq} - h_{rq} R_{pq} gamma_{sp})/2 + (pr<->qs)
x2 = reduce(lib.dot, (h1e_mo, x1, dm1))
# part8
# (g_{ps}\delta_{qr}R_rs + g_{qr}\delta_{ps}) * R_pq)/2 + (pr<->qs)
x2 -= numpy.dot((gpq + gpq.T), x1) * .5
# part2
# (-2Vhf_{sp}\delta_{qr}R_pq - 2Vhf_{qr}\delta_{sp}R_rs)/2 + (pr<->qs)
x2[:ncore] += reduce(numpy.dot,
(x1[:ncore, ncore:], vhf_ca[ncore:])) * 2
# part3
# (-Vhf_{sp}gamma_{qr}R_{pq} - Vhf_{qr}gamma_{sp}R_{rs})/2 + (pr<->qs)
x2[ncore:nocc] += reduce(numpy.dot,
(casdm1, x1[ncore:nocc], eris.vhf_c))
# part1
x2[:, ncore:nocc] += numpy.einsum('purv,rv->pu', hdm2, x1[:,
ncore:nocc])
if ncore > 0:
# part4, part5, part6
# Due to x1_rs [4(pq|sr) + 4(pq|rs) - 2(pr|sq) - 2(ps|rq)] for r>s p>q,
# == -x1_sr [4(pq|sr) + 4(pq|rs) - 2(pr|sq) - 2(ps|rq)] for r>s p>q,
# x2[:,:ncore] += H * x1[:,:ncore] => (becuase x1=-x1.T) =>
# x2[:,:ncore] += -H' * x1[:ncore] => (becuase x2-x2.T) =>
# x2[:ncore] += H' * x1[:ncore]
va, vc = casscf.update_jk_in_ah(mo, x1, casdm1, eris)
x2[ncore:nocc] += va
x2[:ncore, ncore:] += vc
# H_oc
s10 = ci1.dot(ci0) * 2
x2[:, :ncore] += (
(h1e_mo[:, :ncore] + eris.vhf_c[:, :ncore]) * s10 + vhf_a) * 2
x2[:, ncore:nocc] += numpy.dot(
h1e_mo[:, ncore:nocc] + eris.vhf_c[:, ncore:nocc], tdm1)
x2[:, ncore:nocc] += g_dm2
x2 -= s10 * gpq
# (pr<->qs)
x2 = x2 - x2.T
return numpy.hstack((casscf.pack_uniq_var(x2) * 2, hci1 * 2))
def g_update(u, fcivec):
uc = u[:, :ncore].copy()
ua = u[:, ncore:nocc].copy()
rmat = u - numpy.eye(nmo)
ra = rmat[:, ncore:nocc].copy()
mo1 = numpy.dot(mo, u)
mo_c = numpy.dot(mo, uc)
mo_a = numpy.dot(mo, ua)
dm_c = numpy.dot(mo_c, mo_c.T) * 2
fcivec *= 1. / numpy.linalg.norm(fcivec)
casdm1, casdm2 = casscf.fcisolver.make_rdm12(fcivec,
ncas,
nelecas,
link_index=linkstr)
#casscf.with_dep4 = False
#casscf.ci_response_space = 3
#casscf.ci_grad_trust_region = 3
#casdm1, casdm2, gci, fcivec = casscf.update_casdm(mo, u, fcivec, 0, eris, locals())
dm_a = reduce(numpy.dot, (mo_a, casdm1, mo_a.T))
vj, vk = casscf.get_jk(casscf.mol, (dm_c, dm_a))
vhf_c = reduce(numpy.dot, (mo1.T, vj[0] - vk[0] * .5, mo1[:, :nocc]))
vhf_a = reduce(numpy.dot, (mo1.T, vj[1] - vk[1] * .5, mo1[:, :nocc]))
h1e_mo1 = reduce(numpy.dot, (u.T, h1e_mo, u[:, :nocc]))
p1aa = numpy.empty((nmo, ncas, ncas * ncas))
paa1 = numpy.empty((nmo, ncas * ncas, ncas))
aaaa = numpy.empty([ncas] * 4)
for i in range(nmo):
jbuf = eris.ppaa[i]
kbuf = eris.papa[i]
p1aa[i] = lib.dot(ua.T, jbuf.reshape(nmo, -1))
paa1[i] = lib.dot(kbuf.transpose(0, 2, 1).reshape(-1, nmo), ra)
if ncore <= i < nocc:
aaaa[i - ncore] = jbuf[ncore:nocc]
# active space Hamiltonian up to 2nd order
aa11 = lib.dot(ua.T, p1aa.reshape(nmo, -1)).reshape([ncas] * 4)
aa11 = aa11 + aa11.transpose(2, 3, 0, 1) - aaaa
a11a = lib.dot(ra.T, paa1.reshape(nmo, -1)).reshape((ncas, ) * 4)
a11a = a11a + a11a.transpose(1, 0, 2, 3)
a11a = a11a + a11a.transpose(0, 1, 3, 2)
eri_cas_2 = aa11 + a11a
h1cas_2 = h1e_mo1[ncore:nocc, ncore:nocc] + vhf_c[ncore:nocc,
ncore:nocc]
fcivec = fcivec.ravel()
hc0 = fci_matvec(fcivec, h1cas_2, eri_cas_2)
gci = hc0 - fcivec * fcivec.dot(hc0)
g = numpy.zeros_like(h1e_mo)
g[:, :ncore] = (h1e_mo1[:, :ncore] + vhf_c[:, :ncore] +
vhf_a[:, :ncore]) * 2
g[:, ncore:nocc] = numpy.dot(
h1e_mo1[:, ncore:nocc] + vhf_c[:, ncore:nocc], casdm1)
# 0000 + 1000 + 0100 + 0010 + 0001 + 1100 + 1010 + 1001 (missing 0110 + 0101 + 0011)
p1aa = lib.dot(u.T, p1aa.reshape(nmo,
-1)).reshape(nmo, ncas, ncas, ncas)
paa1 = lib.dot(u.T, paa1.reshape(nmo,
-1)).reshape(nmo, ncas, ncas, ncas)
p1aa += paa1
p1aa += paa1.transpose(0, 1, 3, 2)
g[:, ncore:nocc] += numpy.einsum('puwx,wxuv->pv', p1aa, casdm2)
g_orb = casscf.pack_uniq_var(g - g.T)
return numpy.hstack((g_orb * 2, gci * 2))
def get_F(self, k1, k2, use_cov=True, cov_flagging=True):
nchan = self.x[k1].shape[0]
if use_cov:
if not cov_flagging:
F = np.zeros((nchan, nchan), dtype=np.complex)
iC1, iC2 = self.iC(k1), self.iC(k2)
Ctrue1, Ctrue2 = self.Ctrue(k1), self.Ctrue(k2)
else:
# This is for the "effective" matrix s.t. W=MF and p=Mq
F = {}
w1, w2 = self.w[k1], self.w[k2]
m1s = [hash(w1[:, i]) for i in xrange(w1.shape[1])]
m2s = [hash(w2[:, i]) for i in xrange(w2.shape[1])]
for m1, m2 in zip(m1s, m2s):
F[(k1, m1, k2, m2)] = None
for k1, m1, k2, m2 in F.keys():
#for m1 in self._iCt[k1]: # XXX not all m1/m2 pairs may exist in data
# for m2 in self._iCt[k2]:
F[(k1, m1, k2, m2)] = np.zeros((nchan, nchan),
dtype=np.complex)
iCQ1, iCQ2 = {}, {}
for ch in xrange(nchan): # this loop is nchan^3
Q = get_Q(ch, nchan)
iCQ1[ch] = np.dot(
self._iCt[k1][m1],
Q) #C^-1 Q # If ERROR: Compute q_hat first
iCQ2[ch] = np.dot(self._iCt[k2][m2], Q) #C^-1 Q
for i in xrange(nchan): # this loop goes as nchan^4
for j in xrange(nchan):
F[(k1, m1, k2,
m2)][i, j] += np.einsum('ij,ji', iCQ1[i],
iCQ2[j]) #C^-1 Q C^-1 Q
return F
else:
#iC1 = np.linalg.inv(self.C(k1) * np.identity(nchan))
#iC2 = np.linalg.inv(self.C(k2) * np.identity(nchan))
iC1, iC2 = self.I(k1), self.I(k2)
Ctrue1, Ctrue2 = self.I(k1), self.I(
k2) # do this to get the effective F (see below)
F = np.zeros((nchan, nchan), dtype=np.complex)
#Ctrue1, Ctrue2 = self.Ctrue(k1), self.Ctrue(k2)
if False: # This is for the "true" Fisher matrix
CE1, CE2 = {}, {}
for ch in xrange(nchan):
Q = get_Q(ch, nchan)
CE1[ch] = np.dot(Ctrue1,
np.dot(iC1,
np.dot(Q,
iC2))) # C1 Cbar1^-1 Q Cbar2^-1
CE2[ch] = np.dot(Ctrue2,
np.dot(iC2,
np.dot(Q,
iC1))) # C2 Cbar2^-1 Q Cbar1^-1
#CE1[ch] = np.einsum('ab,bc,cd,de', self.Ctrue(k1), iC1, Q, iC2) # slow
#CE2[ch] = np.einsum('ab,bc,cd,de', self.Ctrue(k2), iC2, Q, iC1) # slow
for i in xrange(nchan):
for j in xrange(nchan):
F[i, j] += np.einsum('ij,ji', CE1[i], CE2[j]) # C E C E
else: # This is for the "effective" matrix s.t. W=MF and p=Mq
iCQ1, iCQ2 = {}, {}
for ch in xrange(nchan): # this loop is nchan^3
Q = get_Q(ch, nchan)
iCQ1[ch] = np.dot(iC1, Q) #C^-1 Q
iCQ2[ch] = np.dot(iC2, Q) #C^-1 Q
for i in xrange(nchan): # this loop goes as nchan^4
for j in xrange(nchan):
F[i, j] += np.einsum('ij,ji', iCQ1[i],
iCQ2[j]) #C^-1 Q C^-1 Q
return F
def __init__(self,corner1,corner2,corner3,corner4=None,maskInside=True,coordinates=None,name=None):
"""Generate a box mask with side corner1, corner 2, corner 3 and use corner 4 to define height.
args:
corner1 (3D point): First corner
corner2 (3D point): Second corner
corner2 (3D point): Third corner
kwargs:
corner4 (3D point): Fourth corner used to define height of box. If corner4 is not set corner1, corner2, corner3 are assumed to be maximal positions.
maskInside (bool): If true, points inside is masked otherwise outside (default True)
coordinates (list of str): List containing names of attributes on point object. If None, use x,y (default None)
"""
super(boxMask,self).__init__(coordinates=coordinates,maskInside=maskInside,name=name)
self.corner1 = np.array(corner1,dtype=float)
self.corner2 = np.array(corner2,dtype=float)
self.corner3 = np.array(corner3,dtype=float)
self.corner4 = np.array(corner4,dtype=float)
if corner4 is None: # use provided corners and find bounding box
X,Y,Z = [[np.min(x),np.max(x)] for x in np.array([corner1,corner2,corner3]).T]
c1 = np.array([X[0],Y[0],Z[0]])
c2 = np.array([X[1],Y[0],Z[0]])
c3 = np.array([X[0],Y[1],Z[0]])
c4 = np.array([X[1],Y[1],Z[0]])
c5 = np.array([X[0],Y[0],Z[1]])
c6 = np.array([X[1],Y[0],Z[1]])
c7 = np.array([X[0],Y[1],Z[1]])
c8 = np.array([X[1],Y[1],Z[1]])
self.length = X[1]-X[0]
self.width = Y[1]-Y[0]
self.height = Z[1]-Z[0]
self.planeNormal = np.array([0,0,1])#np.cross(edge1,edge2)
#planeNormal *= 1.0/np.linalg.norm(planeNormal)
#self.rotationMatrix = RotationMatrix3D(0.0,n=None,deg=False)
else:
c1 = np.array(corner1,dtype=float)
c2 = np.array(corner2,dtype=float)
# Start by finding common plane for c1, c2, and c3
edge1 = c2 - c1 # Defined as line between first and second corner
edge2 = corner3 - c1 # Defined as line between first and second corner
self.length = np.linalg.norm(edge1)
edge1 *= 1.0/self.length
edge2 *= 1.0/np.linalg.norm(edge2)
planeNormal = np.cross(edge1,edge2)
self.planeNormal = planeNormal/np.linalg.norm(planeNormal)
cross = np.cross(edge1,self.planeNormal)
width = np.dot(cross,corner3 - c1) # distance from c3 orthogonally on edge1
widthSign = np.sign(width)
self.width = np.abs(width)
# calculate corner point for c3 and c4
c4 = c2+widthSign*cross*self.width
c3 = c4-c2+c1 # Find corresponding c4
# find cube
height = np.dot(self.planeNormal,corner4-c1)
heightSign = np.sign(height)
self.height = np.abs(height)
c5 = c1+heightSign*self.planeNormal*self.height
c6 = c5+c2-c1
c7 = c5+c3-c1
c8 = c5+c4-c1
self.corners = np.array([c1,c2,c3,c4,c5,c6,c7,c8])
self.center = np.array([np.mean(X) for X in self.corners.T]).reshape(1,-1)
# Create rotation matrix.
# rotate planeNormal to lie along x,z plane (https://math.stackexchange.com/questions/180418/calculate-rotation-matrix-to-align-vector-a-to-vector-b-in-3d)
v = np.cross(self.planeNormal,np.array([0.0,1.0,0.0]))
Vs = np.zeros((3,3))
Vs[0,1] = -v[2]
Vs[0,2] = +v[1]
Vs[1,2] = -v[0]
Vs[1,0] = -Vs[0,1]
Vs[0,2] = -Vs[0,2]
Vs[2,1] = -Vs[1,2]
c = np.dot(self.planeNormal,np.array([0.0,1.0,0.0]))
rot1 = np.eye(3)+Vs+np.dot(Vs,Vs)*1.0/(1+c)
c1r1,c2r1,c3r1,c4r1,c5r1,c6r1,c7r1,c8r1 = [np.dot(rot1,(x-self.center).T) for x in [c1,c2,c3,c4,c5,c6,c7,c8]]
# rotate around y into plane using edge 1 (c2-c1)
edge = c2r1-c1r1
theta = np.arctan2(edge[2],edge[0])
rot2 = RotationMatrix3D(theta,n=[0,1,0],deg=False)
self.rotationMatrix = np.einsum('ij,jk->ik',rot2,rot1)
def update_amps(cc, t1, t2, eris):
time0 = time.clock(), time.time()
log = logger.Logger(cc.stdout, cc.verbose)
t1a, t1b = t1
t2aa, t2ab, t2bb = t2
nocca, noccb, nvira, nvirb = t2ab.shape
mo_ea_o = eris.mo_energy[0][:nocca]
mo_ea_v = eris.mo_energy[0][nocca:]
mo_eb_o = eris.mo_energy[1][:noccb]
mo_eb_v = eris.mo_energy[1][noccb:]
fova = eris.focka[:nocca, nocca:]
fovb = eris.fockb[:noccb, noccb:]
u1a = np.zeros_like(t1a)
u1b = np.zeros_like(t1b)
#:eris_vvvv = ao2mo.restore(1, np.asarray(eris.vvvv), nvirb)
#:eris_VVVV = ao2mo.restore(1, np.asarray(eris.VVVV), nvirb)
#:eris_vvVV = _restore(np.asarray(eris.vvVV), nvira, nvirb)
#:u2aa += lib.einsum('ijef,aebf->ijab', tauaa, eris_vvvv) * .5
#:u2bb += lib.einsum('ijef,aebf->ijab', taubb, eris_VVVV) * .5
#:u2ab += lib.einsum('iJeF,aeBF->iJaB', tauab, eris_vvVV)
tauaa, tauab, taubb = make_tau(t2, t1, t1)
u2aa, u2ab, u2bb = cc._add_vvvv(None, (tauaa, tauab, taubb), eris)
u2aa *= .5
u2bb *= .5
Fooa = .5 * lib.einsum('me,ie->mi', fova, t1a)
Foob = .5 * lib.einsum('me,ie->mi', fovb, t1b)
Fvva = -.5 * lib.einsum('me,ma->ae', fova, t1a)
Fvvb = -.5 * lib.einsum('me,ma->ae', fovb, t1b)
Fooa += eris.focka[:nocca, :nocca] - np.diag(mo_ea_o)
Foob += eris.fockb[:noccb, :noccb] - np.diag(mo_eb_o)
Fvva += eris.focka[nocca:, nocca:] - np.diag(mo_ea_v)
Fvvb += eris.fockb[noccb:, noccb:] - np.diag(mo_eb_v)
dtype = u2aa.dtype
wovvo = np.zeros((nocca, nvira, nvira, nocca), dtype=dtype)
wOVVO = np.zeros((noccb, nvirb, nvirb, noccb), dtype=dtype)
woVvO = np.zeros((nocca, nvirb, nvira, noccb), dtype=dtype)
woVVo = np.zeros((nocca, nvirb, nvirb, nocca), dtype=dtype)
wOvVo = np.zeros((noccb, nvira, nvirb, nocca), dtype=dtype)
wOvvO = np.zeros((noccb, nvira, nvira, noccb), dtype=dtype)
mem_now = lib.current_memory()[0]
max_memory = max(0, cc.max_memory - mem_now)
if nvira > 0 and nocca > 0:
blksize = max(ccsd.BLKMIN,
int(max_memory * 1e6 / 8 / (nvira**3 * 3 + 1)))
for p0, p1 in lib.prange(0, nocca, blksize):
ovvv = eris.get_ovvv(slice(p0, p1)) # ovvv = eris.ovvv[p0:p1]
ovvv = ovvv - ovvv.transpose(0, 3, 2, 1)
Fvva += np.einsum('mf,mfae->ae', t1a[p0:p1], ovvv)
wovvo[p0:p1] += lib.einsum('jf,mebf->mbej', t1a, ovvv)
u1a += 0.5 * lib.einsum('mief,meaf->ia', t2aa[p0:p1], ovvv)
u2aa[:, p0:p1] += lib.einsum('ie,mbea->imab', t1a, ovvv.conj())
tmp1aa = lib.einsum('ijef,mebf->ijmb', tauaa, ovvv)
u2aa -= lib.einsum('ijmb,ma->ijab', tmp1aa, t1a[p0:p1] * .5)
ovvv = tmp1aa = None
if nvirb > 0 and noccb > 0:
blksize = max(ccsd.BLKMIN,
int(max_memory * 1e6 / 8 / (nvirb**3 * 3 + 1)))
for p0, p1 in lib.prange(0, noccb, blksize):
OVVV = eris.get_OVVV(slice(p0, p1)) # OVVV = eris.OVVV[p0:p1]
OVVV = OVVV - OVVV.transpose(0, 3, 2, 1)
Fvvb += np.einsum('mf,mfae->ae', t1b[p0:p1], OVVV)
wOVVO[p0:p1] = lib.einsum('jf,mebf->mbej', t1b, OVVV)
u1b += 0.5 * lib.einsum('MIEF,MEAF->IA', t2bb[p0:p1], OVVV)
u2bb[:, p0:p1] += lib.einsum('ie,mbea->imab', t1b, OVVV.conj())
tmp1bb = lib.einsum('ijef,mebf->ijmb', taubb, OVVV)
u2bb -= lib.einsum('ijmb,ma->ijab', tmp1bb, t1b[p0:p1] * .5)
OVVV = tmp1bb = None
if nvirb > 0 and nocca > 0:
blksize = max(ccsd.BLKMIN,
int(max_memory * 1e6 / 8 / (nvira * nvirb**2 * 3 + 1)))
for p0, p1 in lib.prange(0, nocca, blksize):
ovVV = eris.get_ovVV(slice(p0, p1)) # ovVV = eris.ovVV[p0:p1]
Fvvb += np.einsum('mf,mfAE->AE', t1a[p0:p1], ovVV)
woVvO[p0:p1] = lib.einsum('JF,meBF->mBeJ', t1b, ovVV)
woVVo[p0:p1] = lib.einsum('jf,mfBE->mBEj', -t1a, ovVV)
u1b += lib.einsum('mIeF,meAF->IA', t2ab[p0:p1], ovVV)
u2ab[p0:p1] += lib.einsum('IE,maEB->mIaB', t1b, ovVV.conj())
tmp1ab = lib.einsum('iJeF,meBF->iJmB', tauab, ovVV)
u2ab -= lib.einsum('iJmB,ma->iJaB', tmp1ab, t1a[p0:p1])
ovVV = tmp1ab = None
if nvira > 0 and noccb > 0:
blksize = max(ccsd.BLKMIN,
int(max_memory * 1e6 / 8 / (nvirb * nvira**2 * 3 + 1)))
for p0, p1 in lib.prange(0, noccb, blksize):
OVvv = eris.get_OVvv(slice(p0, p1)) # OVvv = eris.OVvv[p0:p1]
Fvva += np.einsum('MF,MFae->ae', t1b[p0:p1], OVvv)
wOvVo[p0:p1] = lib.einsum('jf,MEbf->MbEj', t1a, OVvv)
wOvvO[p0:p1] = lib.einsum('JF,MFbe->MbeJ', -t1b, OVvv)
u1a += lib.einsum('iMfE,MEaf->ia', t2ab[:, p0:p1], OVvv)
u2ab[:, p0:p1] += lib.einsum('ie,MBea->iMaB', t1a, OVvv.conj())
tmp1abba = lib.einsum('iJeF,MFbe->iJbM', tauab, OVvv)
u2ab -= lib.einsum('iJbM,MA->iJbA', tmp1abba, t1b[p0:p1])
OVvv = tmp1abba = None
eris_ovov = np.asarray(eris.ovov)
eris_ovoo = np.asarray(eris.ovoo)
Woooo = lib.einsum('je,nemi->mnij', t1a, eris_ovoo)
Woooo = Woooo - Woooo.transpose(0, 1, 3, 2)
Woooo += np.asarray(eris.oooo).transpose(0, 2, 1, 3)
Woooo += lib.einsum('ijef,menf->mnij', tauaa, eris_ovov) * .5
u2aa += lib.einsum('mnab,mnij->ijab', tauaa, Woooo * .5)
Woooo = tauaa = None
ovoo = eris_ovoo - eris_ovoo.transpose(2, 1, 0, 3)
Fooa += np.einsum('ne,nemi->mi', t1a, ovoo)
u1a += 0.5 * lib.einsum('mnae,meni->ia', t2aa, ovoo)
wovvo += lib.einsum('nb,nemj->mbej', t1a, ovoo)
ovoo = eris_ovoo = None
tilaa = make_tau_aa(t2[0], t1a, t1a, fac=0.5)
ovov = eris_ovov - eris_ovov.transpose(0, 3, 2, 1)
Fvva -= .5 * lib.einsum('mnaf,menf->ae', tilaa, ovov)
Fooa += .5 * lib.einsum('inef,menf->mi', tilaa, ovov)
Fova = np.einsum('nf,menf->me', t1a, ovov)
u2aa += ovov.conj().transpose(0, 2, 1, 3) * .5
wovvo -= 0.5 * lib.einsum('jnfb,menf->mbej', t2aa, ovov)
woVvO += 0.5 * lib.einsum('nJfB,menf->mBeJ', t2ab, ovov)
tmpaa = lib.einsum('jf,menf->mnej', t1a, ovov)
wovvo -= lib.einsum('nb,mnej->mbej', t1a, tmpaa)
eirs_ovov = ovov = tmpaa = tilaa = None
eris_OVOV = np.asarray(eris.OVOV)
eris_OVOO = np.asarray(eris.OVOO)
WOOOO = lib.einsum('je,nemi->mnij', t1b, eris_OVOO)
WOOOO = WOOOO - WOOOO.transpose(0, 1, 3, 2)
WOOOO += np.asarray(eris.OOOO).transpose(0, 2, 1, 3)
WOOOO += lib.einsum('ijef,menf->mnij', taubb, eris_OVOV) * .5
u2bb += lib.einsum('mnab,mnij->ijab', taubb, WOOOO * .5)
WOOOO = taubb = None
OVOO = eris_OVOO - eris_OVOO.transpose(2, 1, 0, 3)
Foob += np.einsum('ne,nemi->mi', t1b, OVOO)
u1b += 0.5 * lib.einsum('mnae,meni->ia', t2bb, OVOO)
wOVVO += lib.einsum('nb,nemj->mbej', t1b, OVOO)
OVOO = eris_OVOO = None
tilbb = make_tau_aa(t2[2], t1b, t1b, fac=0.5)
OVOV = eris_OVOV - eris_OVOV.transpose(0, 3, 2, 1)
Fvvb -= .5 * lib.einsum('MNAF,MENF->AE', tilbb, OVOV)
Foob += .5 * lib.einsum('inef,menf->mi', tilbb, OVOV)
Fovb = np.einsum('nf,menf->me', t1b, OVOV)
u2bb += OVOV.conj().transpose(0, 2, 1, 3) * .5
wOVVO -= 0.5 * lib.einsum('jnfb,menf->mbej', t2bb, OVOV)
wOvVo += 0.5 * lib.einsum('jNbF,MENF->MbEj', t2ab, OVOV)
tmpbb = lib.einsum('jf,menf->mnej', t1b, OVOV)
wOVVO -= lib.einsum('nb,mnej->mbej', t1b, tmpbb)
eris_OVOV = OVOV = tmpbb = tilbb = None
eris_OVoo = np.asarray(eris.OVoo)
eris_ovOO = np.asarray(eris.ovOO)
Fooa += np.einsum('NE,NEmi->mi', t1b, eris_OVoo)
u1a -= lib.einsum('nMaE,MEni->ia', t2ab, eris_OVoo)
wOvVo -= lib.einsum('nb,MEnj->MbEj', t1a, eris_OVoo)
woVVo += lib.einsum('NB,NEmj->mBEj', t1b, eris_OVoo)
Foob += np.einsum('ne,neMI->MI', t1a, eris_ovOO)
u1b -= lib.einsum('mNeA,meNI->IA', t2ab, eris_ovOO)
woVvO -= lib.einsum('NB,meNJ->mBeJ', t1b, eris_ovOO)
wOvvO += lib.einsum('nb,neMJ->MbeJ', t1a, eris_ovOO)
WoOoO = lib.einsum('JE,NEmi->mNiJ', t1b, eris_OVoo)
WoOoO += lib.einsum('je,neMI->nMjI', t1a, eris_ovOO)
WoOoO += np.asarray(eris.ooOO).transpose(0, 2, 1, 3)
eris_OVoo = eris_ovOO = None
eris_ovOV = np.asarray(eris.ovOV)
WoOoO += lib.einsum('iJeF,meNF->mNiJ', tauab, eris_ovOV)
u2ab += lib.einsum('mNaB,mNiJ->iJaB', tauab, WoOoO)
WoOoO = None
tilab = make_tau_ab(t2[1], t1, t1, fac=0.5)
Fvva -= lib.einsum('mNaF,meNF->ae', tilab, eris_ovOV)
Fvvb -= lib.einsum('nMfA,nfME->AE', tilab, eris_ovOV)
Fooa += lib.einsum('iNeF,meNF->mi', tilab, eris_ovOV)
Foob += lib.einsum('nIfE,nfME->MI', tilab, eris_ovOV)
Fova += np.einsum('NF,meNF->me', t1b, eris_ovOV)
Fovb += np.einsum('nf,nfME->ME', t1a, eris_ovOV)
u2ab += eris_ovOV.conj().transpose(0, 2, 1, 3)
wovvo += 0.5 * lib.einsum('jNbF,meNF->mbej', t2ab, eris_ovOV)
wOVVO += 0.5 * lib.einsum('nJfB,nfME->MBEJ', t2ab, eris_ovOV)
wOvVo -= 0.5 * lib.einsum('jnfb,nfME->MbEj', t2aa, eris_ovOV)
woVvO -= 0.5 * lib.einsum('JNFB,meNF->mBeJ', t2bb, eris_ovOV)
woVVo += 0.5 * lib.einsum('jNfB,mfNE->mBEj', t2ab, eris_ovOV)
wOvvO += 0.5 * lib.einsum('nJbF,neMF->MbeJ', t2ab, eris_ovOV)
tmpabab = lib.einsum('JF,meNF->mNeJ', t1b, eris_ovOV)
tmpbaba = lib.einsum('jf,nfME->MnEj', t1a, eris_ovOV)
woVvO -= lib.einsum('NB,mNeJ->mBeJ', t1b, tmpabab)
wOvVo -= lib.einsum('nb,MnEj->MbEj', t1a, tmpbaba)
woVVo += lib.einsum('NB,NmEj->mBEj', t1b, tmpbaba)
wOvvO += lib.einsum('nb,nMeJ->MbeJ', t1a, tmpabab)
tmpabab = tmpbaba = tilab = None
Fova += fova
Fovb += fovb
u1a += fova.conj()
u1a += np.einsum('ie,ae->ia', t1a, Fvva)
u1a -= np.einsum('ma,mi->ia', t1a, Fooa)
u1a -= np.einsum('imea,me->ia', t2aa, Fova)
u1a += np.einsum('iMaE,ME->ia', t2ab, Fovb)
u1b += fovb.conj()
u1b += np.einsum('ie,ae->ia', t1b, Fvvb)
u1b -= np.einsum('ma,mi->ia', t1b, Foob)
u1b -= np.einsum('imea,me->ia', t2bb, Fovb)
u1b += np.einsum('mIeA,me->IA', t2ab, Fova)
eris_oovv = np.asarray(eris.oovv)
eris_ovvo = np.asarray(eris.ovvo)
wovvo -= eris_oovv.transpose(0, 2, 3, 1)
wovvo += eris_ovvo.transpose(0, 2, 1, 3)
oovv = eris_oovv - eris_ovvo.transpose(0, 3, 2, 1)
u1a -= np.einsum('nf,niaf->ia', t1a, oovv)
tmp1aa = lib.einsum('ie,mjbe->mbij', t1a, oovv)
u2aa += 2 * lib.einsum('ma,mbij->ijab', t1a, tmp1aa)
eris_ovvo = eris_oovv = oovv = tmp1aa = None
eris_OOVV = np.asarray(eris.OOVV)
eris_OVVO = np.asarray(eris.OVVO)
wOVVO -= eris_OOVV.transpose(0, 2, 3, 1)
wOVVO += eris_OVVO.transpose(0, 2, 1, 3)
OOVV = eris_OOVV - eris_OVVO.transpose(0, 3, 2, 1)
u1b -= np.einsum('nf,niaf->ia', t1b, OOVV)
tmp1bb = lib.einsum('ie,mjbe->mbij', t1b, OOVV)
u2bb += 2 * lib.einsum('ma,mbij->ijab', t1b, tmp1bb)
eris_OVVO = eris_OOVV = OOVV = None
eris_ooVV = np.asarray(eris.ooVV)
eris_ovVO = np.asarray(eris.ovVO)
woVVo -= eris_ooVV.transpose(0, 2, 3, 1)
woVvO += eris_ovVO.transpose(0, 2, 1, 3)
u1b += np.einsum('nf,nfAI->IA', t1a, eris_ovVO)
tmp1ab = lib.einsum('ie,meBJ->mBiJ', t1a, eris_ovVO)
tmp1ab += lib.einsum('IE,mjBE->mBjI', t1b, eris_ooVV)
u2ab -= lib.einsum('ma,mBiJ->iJaB', t1a, tmp1ab)
eris_ooVV = eris_ovVo = tmp1ab = None
eris_OOvv = np.asarray(eris.OOvv)
eris_OVvo = np.asarray(eris.OVvo)
wOvvO -= eris_OOvv.transpose(0, 2, 3, 1)
wOvVo += eris_OVvo.transpose(0, 2, 1, 3)
u1a += np.einsum('NF,NFai->ia', t1b, eris_OVvo)
tmp1ba = lib.einsum('IE,MEbj->MbIj', t1b, eris_OVvo)
tmp1ba += lib.einsum('ie,MJbe->MbJi', t1a, eris_OOvv)
u2ab -= lib.einsum('MA,MbIj->jIbA', t1b, tmp1ba)
eris_OOvv = eris_OVvO = tmp1ba = None
u2aa += 2 * lib.einsum('imae,mbej->ijab', t2aa, wovvo)
u2aa += 2 * lib.einsum('iMaE,MbEj->ijab', t2ab, wOvVo)
u2bb += 2 * lib.einsum('imae,mbej->ijab', t2bb, wOVVO)
u2bb += 2 * lib.einsum('mIeA,mBeJ->IJAB', t2ab, woVvO)
u2ab += lib.einsum('imae,mBeJ->iJaB', t2aa, woVvO)
u2ab += lib.einsum('iMaE,MBEJ->iJaB', t2ab, wOVVO)
u2ab += lib.einsum('iMeA,MbeJ->iJbA', t2ab, wOvvO)
u2ab += lib.einsum('IMAE,MbEj->jIbA', t2bb, wOvVo)
u2ab += lib.einsum('mIeA,mbej->jIbA', t2ab, wovvo)
u2ab += lib.einsum('mIaE,mBEj->jIaB', t2ab, woVVo)
wovvo = wOVVO = woVvO = wOvVo = woVVo = wOvvO = None
Ftmpa = Fvva - .5 * lib.einsum('mb,me->be', t1a, Fova)
Ftmpb = Fvvb - .5 * lib.einsum('mb,me->be', t1b, Fovb)
u2aa += lib.einsum('ijae,be->ijab', t2aa, Ftmpa)
u2bb += lib.einsum('ijae,be->ijab', t2bb, Ftmpb)
u2ab += lib.einsum('iJaE,BE->iJaB', t2ab, Ftmpb)
u2ab += lib.einsum('iJeA,be->iJbA', t2ab, Ftmpa)
Ftmpa = Fooa + 0.5 * lib.einsum('je,me->mj', t1a, Fova)
Ftmpb = Foob + 0.5 * lib.einsum('je,me->mj', t1b, Fovb)
u2aa -= lib.einsum('imab,mj->ijab', t2aa, Ftmpa)
u2bb -= lib.einsum('imab,mj->ijab', t2bb, Ftmpb)
u2ab -= lib.einsum('iMaB,MJ->iJaB', t2ab, Ftmpb)
u2ab -= lib.einsum('mIaB,mj->jIaB', t2ab, Ftmpa)
eris_ovoo = np.asarray(eris.ovoo).conj()
eris_OVOO = np.asarray(eris.OVOO).conj()
eris_OVoo = np.asarray(eris.OVoo).conj()
eris_ovOO = np.asarray(eris.ovOO).conj()
ovoo = eris_ovoo - eris_ovoo.transpose(2, 1, 0, 3)
OVOO = eris_OVOO - eris_OVOO.transpose(2, 1, 0, 3)
u2aa -= lib.einsum('ma,jbim->ijab', t1a, ovoo)
u2bb -= lib.einsum('ma,jbim->ijab', t1b, OVOO)
u2ab -= lib.einsum('ma,JBim->iJaB', t1a, eris_OVoo)
u2ab -= lib.einsum('MA,ibJM->iJbA', t1b, eris_ovOO)
eris_ovoo = eris_OVoo = eris_OVOO = eris_ovOO = None
u2aa *= .5
u2bb *= .5
u2aa = u2aa - u2aa.transpose(0, 1, 3, 2)
u2aa = u2aa - u2aa.transpose(1, 0, 2, 3)
u2bb = u2bb - u2bb.transpose(0, 1, 3, 2)
u2bb = u2bb - u2bb.transpose(1, 0, 2, 3)
eia_a = lib.direct_sum('i-a->ia', mo_ea_o, mo_ea_v)
eia_b = lib.direct_sum('i-a->ia', mo_eb_o, mo_eb_v)
u1a /= eia_a
u1b /= eia_b
u2aa /= lib.direct_sum('ia+jb->ijab', eia_a, eia_a)
u2ab /= lib.direct_sum('ia+jb->ijab', eia_a, eia_b)
u2bb /= lib.direct_sum('ia+jb->ijab', eia_b, eia_b)
time0 = log.timer_debug1('update t1 t2', *time0)
t1new = u1a, u1b
t2new = u2aa, u2ab, u2bb
return t1new, t2new
def DFTExcitedState(mol, func, orbitals, **kwargs):
"""
Perform unrestrictred Kohn-Sham excited state calculation
"""
psi4.core.print_out("\nEntering Excited State Kohn-Sham:\n" + 33 * "=" +
"\n\n")
maxiter = int(psi4.core.get_local_option("PSIXAS", "MAXITER"))
E_conv = 1.0E-8
D_conv = 1.0E-6
"""
STEP 1: Read in ground state orbitals or restart from previous
"""
prefix = psi4.core.get_local_option("PSIXAS", "PREFIX")
if (os.path.isfile(prefix + "_exorbs.npz") and False):
psi4.core.print_out("Restarting Calculation\n")
Ca = np.load(prefix + "_exorbs.npz")["Ca"]
Cb = np.load(prefix + "_exorbs.npz")["Cb"]
else:
Ca = np.load(prefix + "_gsorbs.npz")["Ca"]
Cb = np.load(prefix + "_gsorbs.npz")["Cb"]
"""
Grep the coefficients for later overlap
"""
for i in orbitals:
if i["spin"] == "b":
i["C"] = Cb[:, i["orb"]]
elif i["spin"] == "a":
i["C"] = Ca[:, i["orb"]]
else:
raise Exception("Orbital has non a/b spin!")
wfn = psi4.core.Wavefunction.build(mol,
psi4.core.get_global_option('BASIS'))
aux = psi4.core.BasisSet.build(mol, "DF_BASIS_SCF", "", "JKFIT",
psi4.core.get_global_option('BASIS'))
#psi4.core.be_quiet()
mints = psi4.core.MintsHelper(wfn.basisset())
S = np.asarray(mints.ao_overlap())
T = np.asarray(mints.ao_kinetic())
V = np.asarray(mints.ao_potential())
H = np.zeros((mints.nbf(), mints.nbf()))
H = T + V
A = mints.ao_overlap()
A.power(-0.5, 1.e-16)
A = np.asarray(A)
Enuc = mol.nuclear_repulsion_energy()
Eold = 0.0
SCF_E = 100.0
nbf = wfn.nso()
nalpha = wfn.nalpha()
nbeta = wfn.nbeta()
Va = psi4.core.Matrix(nbf, nbf)
Vb = psi4.core.Matrix(nbf, nbf)
sup = psi4.driver.dft.build_superfunctional(func, False)[0]
sup.set_deriv(2)
sup.allocate()
Vpot = psi4.core.VBase.build(wfn.basisset(), sup, "UV")
Vpot.initialize()
#This object is needed to write out a molden file later
uhf = psi4.core.UHF(wfn, sup)
psi4.core.reopen_outfile()
"""
Form initial denisty
"""
occa = np.zeros(Ca.shape[0])
occb = np.zeros(Cb.shape[0])
occa[:nalpha] = 1
occb[:nbeta] = 1
Cocca = psi4.core.Matrix(nbf, nbf)
Coccb = psi4.core.Matrix(nbf, nbf)
Cocca.np[:] = Ca
Coccb.np[:] = Cb
for i in orbitals:
if i["spin"] == "b":
"""
Check if this is still the largest overlap
"""
ovl = np.abs(np.einsum('m,nj,mn->j', i["C"], Coccb, S))
if i["orb"] != np.argmax(ovl):
print("index changed from {:d} to {:d}".format(
i["orb"], np.argmax(ovl)))
i["orb"] = np.argmax(ovl)
"""
Set occupation and overlap
"""
i["ovl"] = np.max(ovl)
occb[i["orb"]] = i["occ"]
elif i["spin"] == "a":
"""
Check if this is still the largest overlap
"""
ovl = np.abs(np.einsum('m,nj,mn->j', i["C"], Cocca, S))
if i["orb"] != np.argmax(ovl):
print("index changed from {:d} to {:d}".format(
i["orb"], np.argmax(ovl)))
i["orb"] = np.argmax(ovl)
"""
Set occupation and overlap
"""
i["ovl"] = np.max(ovl)
occa[i["orb"]] = i["occ"]
for i in range(nbf):
Cocca.np[:, i] *= np.sqrt(occa[i])
Coccb.np[:, i] *= np.sqrt(occb[i])
Da = Cocca.np @ Cocca.np.T
Db = Coccb.np @ Coccb.np.T
jk = psi4.core.JK.build(wfn.basisset(), aux, "MEM_DF")
glob_mem = psi4.core.get_memory() / 8
jk.set_memory(int(glob_mem * 0.6))
jk.initialize()
jk.C_left_add(Cocca)
jk.C_left_add(Coccb)
Da_m = psi4.core.Matrix(nbf, nbf)
Db_m = psi4.core.Matrix(nbf, nbf)
diisa = DIIS_helper()
diisb = DIIS_helper()
gamma = psi4.core.get_local_option("PSIXAS", "DAMP")
diis_eps = psi4.core.get_local_option("PSIXAS", "DIIS_EPS")
vshift = psi4.core.get_local_option("PSIXAS", "VSHIFT")
psi4.core.print_out(
"\nStarting SCF:\n" + 13 * "=" +
"\n\n{:>10} {:4.2f}\n{:>10} {:4.2f}\n{:>10} {:4.2f}\n\n".format(
"DAMP:", gamma, "DIIS_EPS:", diis_eps, "VSHIFT:", vshift))
psi4.core.print_out("\nInitial orbital occupation pattern:\n\n")
psi4.core.print_out("Index|Spin|Occ|Ovl|Freeze\n" + 25 * "-")
for i in orbitals:
psi4.core.print_out("\n{:^5}|{:^4}|{:^3}|{:^3}|{:^6}".format(
i["orb"], i["spin"], i["occ"], 'Yes' if i["DoOvl"] else 'No',
'Yes' if i["frz"] else 'No'))
psi4.core.print_out("\n\n")
psi4.core.print_out(
("{:^3} {:^14} {:^14} {:^4} {:^4} | {:^" + str(len(orbitals) * 5) +
"}| {:^4} {:^5}\n").format("#IT", "Escf", "dEscf", "na", "nb", "OVL",
"MIX", "Time"))
psi4.core.print_out("=" * 80 + "\n")
myTimer = Timer()
MIXMODE = "DAMP"
for SCF_ITER in range(1, maxiter + 1):
myTimer.addStart("SCF")
myTimer.addStart("JK")
jk.compute()
myTimer.addEnd("JK")
myTimer.addStart("buildFock")
Da_m.np[:] = Da
Db_m.np[:] = Db
Vpot.set_D([Da_m, Db_m])
myTimer.addStart("compV")
Vpot.compute_V([Va, Vb])
myTimer.addEnd("compV")
Ja = np.asarray(jk.J()[0])
Jb = np.asarray(jk.J()[1])
Ka = np.asarray(jk.K()[0])
Kb = np.asarray(jk.K()[1])
if SCF_ITER > 1:
FaOld = np.copy(Fa)
FbOld = np.copy(Fb)
Fa = H + (Ja + Jb) - Vpot.functional().x_alpha() * Ka + Va
Fb = H + (Ja + Jb) - Vpot.functional().x_alpha() * Kb + Vb
myTimer.addEnd("buildFock")
"""
Fock matrix constructed: freeze orbitals if needed
"""
myTimer.addStart("Freeze")
FMOa = Ca.T @ Fa @ Ca
FMOb = Cb.T @ Fb @ Cb
CaInv = np.linalg.inv(Ca)
CbInv = np.linalg.inv(Cb)
for i in orbitals:
if i['frz'] == True:
if i["spin"] == "b":
idx = i["orb"]
FMOb[idx, :idx] = 0.0
FMOb[idx, (idx + 1):] = 0.0
FMOb[:idx, idx] = 0.0
FMOb[(idx + 1):, idx] = 0.0
elif i["spin"] == "a":
FMOa[idx, :idx] = 0.0
FMOa[idx, (idx + 1):] = 0.0
FMOa[:idx, idx] = 0.0
FMOa[(idx + 1):, idx] = 0.0
"""
VSHIFT
"""
idxs = [c for c, x in enumerate(occa) if (x == 0.0) and (c >= nalpha)]
FMOa[idxs, idxs] += vshift
idxs = [c for c, x in enumerate(occb) if (x == 0.0) and (c >= nbeta)]
FMOb[idxs, idxs] += vshift
Fa = CaInv.T @ FMOa @ CaInv
Fb = CbInv.T @ FMOb @ CbInv
myTimer.addEnd("Freeze")
"""
END FREEZE
"""
"""
DIIS/MIXING
"""
myTimer.addStart("MIX")
diisa_e = Fa.dot(Da).dot(S) - S.dot(Da).dot(Fa)
diisa_e = (A.T).dot(diisa_e).dot(A)
diisa.add(Fa, diisa_e)
diisb_e = Fb.dot(Db).dot(S) - S.dot(Db).dot(Fb)
diisb_e = (A.T).dot(diisb_e).dot(A)
diisb.add(Fb, diisb_e)
if (MIXMODE == "DIIS") and (SCF_ITER > 1):
# Extrapolate alpha & beta Fock matrices separately
Fa = diisa.extrapolate()
Fb = diisb.extrapolate()
elif (MIXMODE == "DAMP") and (SCF_ITER > 1):
# Use Damping to obtain the new Fock matrices
Fa = (1 - gamma) * np.copy(Fa) + (gamma) * FaOld
Fb = (1 - gamma) * np.copy(Fb) + (gamma) * FbOld
myTimer.addEnd("MIX")
"""
END DIIS/MIXING
"""
"""
CALC energy
"""
myTimer.addStart("calcE")
one_electron_E = np.sum(Da * H)
one_electron_E += np.sum(Db * H)
coulomb_E = np.sum(Da * (Ja + Jb))
coulomb_E += np.sum(Db * (Ja + Jb))
alpha = Vpot.functional().x_alpha()
exchange_E = 0.0
exchange_E -= alpha * np.sum(Da * Ka)
exchange_E -= alpha * np.sum(Db * Kb)
XC_E = Vpot.quadrature_values()["FUNCTIONAL"]
SCF_E = 0.0
SCF_E += Enuc
SCF_E += one_electron_E
SCF_E += 0.5 * coulomb_E
SCF_E += 0.5 * exchange_E
SCF_E += XC_E
myTimer.addEnd("calcE")
# Diagonalize Fock matrix
myTimer.addStart("Diag")
Ca, epsa = diag_H(Fa, A)
Cb, epsb = diag_H(Fb, A)
myTimer.addEnd("Diag")
"""
New orbitals obtained set occupation numbers
"""
myTimer.addStart("SetOcc")
Cocca.np[:] = Ca
Coccb.np[:] = Cb
occa[:] = 0.0
occb[:] = 0.0
occa[:nalpha] = 1.0 #standard aufbau principle occupation
occb[:nbeta] = 1.0
for i in orbitals:
if i["spin"] == "b":
"""
Overlap
"""
#calculate the Overlapp with all other orbitals
ovl = np.abs(np.einsum('m,nj,mn->j', i["C"], Coccb, S))
#User wants to switch the index if higher overlap is found
if i["DoOvl"] == True:
if i["orb"] != np.argmax(ovl):
i["orb"] = np.argmax(ovl)
i["ovl"] = np.max(ovl)
else:
#just calculate the overlap to assess the character
i["ovl"] = ovl[i["orb"]]
#Modify the occupation vector
occb[i["orb"]] = i["occ"]
elif i["spin"] == "a":
"""
Check if this is still the largest overlap
"""
ovl = np.abs(np.einsum('m,nj,mn->j', i["C"], Cocca, S))
if i["DoOvl"] == True:
if i["orb"] != np.argmax(ovl):
i["orb"] = np.argmax(
ovl) # set index to the highest overlap
i["ovl"] = np.max(ovl)
else:
i["ovl"] = ovl[i["orb"]]
#Modify the occupation vector
occa[i["orb"]] = i["occ"]
for i in range(nbf):
Cocca.np[:, i] *= np.sqrt(occa[i])
Coccb.np[:, i] *= np.sqrt(occb[i])
Da = Cocca.np @ Cocca.np.T
Db = Coccb.np @ Coccb.np.T
myTimer.addEnd("SetOcc")
myTimer.addEnd("SCF")
psi4.core.print_out(
("{:3d} {:14.8f} {:14.8f} {:4.1f} {:4.1f} | " +
"{:4.2f} " * len(orbitals) + "| {:^4} {:5.2f} \n").format(
SCF_ITER, SCF_E, (SCF_E - Eold), np.sum(Da * S),
np.sum(Db * S), *[x["ovl"] for x in orbitals], MIXMODE,
myTimer.getTime("SCF")))
psi4.core.flush_outfile()
myTimer.printAlltoFile("timers.ksex")
if (abs(SCF_E - Eold) < diis_eps):
MIXMODE = "DIIS"
else:
MIXMODE = "DAMP"
if (abs(SCF_E - Eold) < E_conv):
if (vshift != 0.0):
psi4.core.print_out(
"Converged but Vshift was on... removing Vshift..\n")
vshift = 0.0
else:
break
Eold = SCF_E
if SCF_ITER == maxiter:
psi4.core.clean()
raise Exception("Maximum number of SCF cycles exceeded.")
psi4.core.print_out("\n\n{:>20} {:12.8f} [Ha] \n".format(
"FINAL EX SCF ENERGY:", SCF_E))
gsE = psi4.core.scalar_variable('GS ENERGY')
if gsE != 0.0:
psi4.core.print_out("{:>20} {:12.8f} [Ha] \n".format(
"EXCITATION ENERGY:", SCF_E - gsE))
psi4.core.print_out("{:>20} {:12.8f} [eV] \n\n".format(
"EXCITATION ENERGY:", (SCF_E - gsE) * 27.211385))
psi4.core.print_out("\nFinal orbital occupation pattern:\n\n")
psi4.core.print_out("Index|Spin|Occ|Ovl|Freeze\n" + 25 * "-")
for i in orbitals:
if i["spin"] == "a":
occ = occa[i["orb"]]
else:
occ = occb[i["orb"]]
psi4.core.print_out("\n{:^5}|{:^4}|{:^3}|{:^3}|{:^6}".format(
i["orb"], i["spin"], occ, 'Yes' if i["DoOvl"] else 'No',
'Yes' if i["frz"] else 'No'))
psi4.core.print_out("\n\n")
OCCA = psi4.core.Vector("Alpha occupation", nbf)
OCCB = psi4.core.Vector("Beta occupation", nbf)
OCCA.np[:] = occa
OCCB.np[:] = occb
uhf.Ca().np[:] = Ca
uhf.Cb().np[:] = Cb
uhf.epsilon_a().np[:] = epsa
uhf.epsilon_b().np[:] = epsb
uhf.occupation_a().np[:] = occa
uhf.occupation_b().np[:] = occb
uhf.epsilon_a().print_out()
uhf.epsilon_b().print_out()
OCCA.print_out()
OCCB.print_out()
mw = psi4.core.MoldenWriter(uhf)
mw.write(prefix + '_ex.molden', uhf.Ca(), uhf.Cb(), uhf.epsilon_a(),
uhf.epsilon_b(), OCCA, OCCB, True)
psi4.core.print_out("\n\n Moldenfile written\n")
np.savez(prefix + '_exorbs',
Ca=Ca,
Cb=Cb,
occa=occa,
occb=occb,
epsa=epsa,
epsb=epsb,
orbitals=orbitals)
psi4.core.set_variable('CURRENT ENERGY', SCF_E)
bas = wfn.basisset()
with open("GENBAS", "w") as genbas:
genbas.write(bas.genbas())
maxl = bas.max_am()
map = []
lmap = [[0], [1, 2, 0], [0, 4, 1, 3, 2], [1, 2, 0, 5, 4, 6, 3],
[0, 4, 1, 7, 6, 3, 8, 5, 2], [1, 2, 3, 5, 8, 10, 7, 6, 0, 9, 4],
[11, 4, 9, 7, 10, 3, 12, 5, 8, 0, 6, 2, 1]]
for i in range(mol.natom()):
for l in range(maxl + 1):
for li in range(2 * l + 1):
for j in range(bas.nshell()):
s = bas.shell(j)
if bas.shell_to_center(j) == i and s.am == l:
k = bas.shell_to_basis_function(j)
if l > 6:
m = l - li // 2
k += 2 * m + li % 2
else:
k += lmap[l][li]
map.append(k)
assert (len(map) == mints.nbf())
assert (sorted(map) == list(range(mints.nbf())))
with open("OLDMOS", "w") as oldmos:
for C in (Ca, Cb):
for j0 in range(0, mints.nbf(), 4):
for i in range(mints.nbf()):
for j in range(j0, min(mints.nbf(), j0 + 4)):
oldmos.write("%30.20e" % (C[map[i]][j]))
oldmos.write('\n')
open("JFSGUESS", "w").close()
def image_of_modules(image):
return np.sqrt(np.einsum('ijk,ijk->ij', image, image))
axes_b = [axes_b[x] for x in perm]
res = tensordot(a, b, axes=(tuple(axes_a[::-1]), tuple(axes_b[::-1])))
if rkey != rkey_default:
perm = [rkey_default.index(k) for k in rkey]
res = res.transpose(*perm)
return res
if __name__ == '__main__':
keys = [
'ibjc,jia->cab',
'cab,icab->i',
'ibjc,jia->cab',
'cab,icab->i',
'i,ibac->cab',
'cab,jcib->ija',
'blk,lkib->i',
]
sizes = dict(i=5, j=5, k=5, l=5, a=5, b=5, c=5, d=5)
for key in keys:
a = np.random.random([sizes[x] for x in key.split(',')[0]])
b = np.random.random(
[sizes[x] for x in key.split(',')[1].split('->')[0]])
resa = np.einsum(key, a, b)
resb = einsum(key, a, b)
print('%16s %8s' % (key, np.allclose(resa, resb)))
class TestEvaluation(TestCase):
# coefficients of 1 + 2*x + 3*x**2
c1d = np.array([4., 2., 3.])
c2d = np.einsum('i,j->ij', c1d, c1d)
c3d = np.einsum('i,j,k->ijk', c1d, c1d, c1d)
# some random values in [-1, 1)
x = np.random.random((3, 5)) * 2 - 1
y = polyval(x, [1., 2., 3.])
def test_hermeval(self):
#check empty input
assert_equal(herme.hermeval([], [1]).size, 0)
#check normal input)
x = np.linspace(-1, 1)
y = [polyval(x, c) for c in Helist]
for i in range(10):
msg = "At i=%d" % i
ser = np.zeros
tgt = y[i]
res = herme.hermeval(x, [0] * i + [1])
assert_almost_equal(res, tgt, err_msg=msg)
#check that shape is preserved
for i in range(3):
dims = [2] * i
x = np.zeros(dims)
assert_equal(herme.hermeval(x, [1]).shape, dims)
assert_equal(herme.hermeval(x, [1, 0]).shape, dims)
assert_equal(herme.hermeval(x, [1, 0, 0]).shape, dims)
def test_hermeval2d(self):
x1, x2, x3 = self.x
y1, y2, y3 = self.y
#test exceptions
assert_raises(ValueError, herme.hermeval2d, x1, x2[:2], self.c2d)
#test values
tgt = y1 * y2
res = herme.hermeval2d(x1, x2, self.c2d)
assert_almost_equal(res, tgt)
#test shape
z = np.ones((2, 3))
res = herme.hermeval2d(z, z, self.c2d)
assert_(res.shape == (2, 3))
def test_hermeval3d(self):
x1, x2, x3 = self.x
y1, y2, y3 = self.y
#test exceptions
assert_raises(ValueError, herme.hermeval3d, x1, x2, x3[:2], self.c3d)
#test values
tgt = y1 * y2 * y3
res = herme.hermeval3d(x1, x2, x3, self.c3d)
assert_almost_equal(res, tgt)
#test shape
z = np.ones((2, 3))
res = herme.hermeval3d(z, z, z, self.c3d)
assert_(res.shape == (2, 3))
def test_hermegrid2d(self):
x1, x2, x3 = self.x
y1, y2, y3 = self.y
#test values
tgt = np.einsum('i,j->ij', y1, y2)
res = herme.hermegrid2d(x1, x2, self.c2d)
assert_almost_equal(res, tgt)
#test shape
z = np.ones((2, 3))
res = herme.hermegrid2d(z, z, self.c2d)
assert_(res.shape == (2, 3) * 2)
def test_hermegrid3d(self):
x1, x2, x3 = self.x
y1, y2, y3 = self.y
#test values
tgt = np.einsum('i,j,k->ijk', y1, y2, y3)
res = herme.hermegrid3d(x1, x2, x3, self.c3d)
assert_almost_equal(res, tgt)
#test shape
z = np.ones((2, 3))
res = herme.hermegrid3d(z, z, z, self.c3d)
assert_(res.shape == (2, 3) * 3)
def find_max_diff_l2(spectra, from_spectrum):
diff = spectra - from_spectrum
distances = np.einsum('ij,ij->i', diff, diff, optimize='optimal')
index_of_max = np.argmax(distances)
return spectra[index_of_max]
def gls_fit(ramp_data,
input_var_data,
prev_fit_data,
prev_slope_data,
readnoise,
gain,
frame_time,
group_time,
nframes_used,
num_cr,
cr_flagged_2d,
saturated_data,
use_extra_terms=True):
"""Generalized least squares linear fit.
It is assumed that every input pixel has num_cr cosmic-ray hits
somewhere within the ramp. This function should be called separately
for different values of num_cr.
Parameters
----------
ramp_data: 2-D ndarray; indices: group, pixel number
The ramp data for one of the integrations in an exposure. This
may be a subset in detector coordinates, but covering all groups.
The shape is (ngroups, nz), where ngroups is the length of the
ramp, and nz is the number of pixels in the current subset.
input_var_data: 2-D ndarray, shape (ngroups, nz)
The square of the input ERR array, matching ramp_data.
prev_fit_data: 2-D ndarray, shape (ngroups, nz)
The fit to ramp_data, based on applying the values of intercept,
slope, and cosmic-ray amplitudes that were determined in a previous
call to gls_fit. This array is only used for setting up the
covariance matrix.
prev_slope_data: 1-D ndarray, length nz.
An estimate (e.g. from a previous iteration) of the slope at each
pixel, in electrons per second.
readnoise: 1-D ndarray, length nz.
The read noise in electrons at each detector pixel.
gain: 1-D ndarray, shape (nz,)
The analog-to-digital gain (electrons per dn) at each detector
pixel.
frame_time: float
The time to read one frame, in seconds (e.g. 10.6 s).
group_time: float
Time increment between groups, in seconds.
nframes_used: int
Number of frames that were averaged together to make a group.
Note that this value does not include the number (if any) of
skipped frames.
num_cr: int
The number of cosmic rays that will be handled. All pixels in the
current set (ramp_data) are assumed to have this many cosmic ray
hits somewhere within the ramp.
cr_flagged_2d: 2-D ndarray, shape (ngroups, nz)
The values should be 0 or 1; 1 indicates that a cosmic ray was
detected (by another step) at that point.
saturated_data: 2-D ndarray, shape (ngroups, nz)
Normal values are zero; the value will be a huge number for
saturated pixels. This will be added to the main diagonal of the
inverse of the weight matrix to greatly reduce the weight for
saturated pixels.
use_extra_terms: bool
True if we should include Massimo Robberto's terms in the inverse
weight matrix.
See JWST-STScI-003193.pdf
Returns
-------
tuple: (result2d, variances)
result2d is a 2-D ndarray; shape (nz, 2 + num_cr)
The computed values of intercept, slope, and cosmic-ray amplitudes
(there will be num_cr cosmic-ray amplitudes) for each of the nz
pixels.
variances is a 2-D ndarray; shape (nz, 2 + num_cr)
The variance for the intercept, slope, and for the amplitude of
each cosmic ray that was detected.
"""
M = float(nframes_used)
ngroups = ramp_data.shape[0]
nz = ramp_data.shape[1]
num_cr = int(num_cr)
# x is an array (length nz) of matrices, each of which is the
# independent variable of a linear equation. Each such matrix
# has ngroups rows and 2 + num_cr columns. The first column is set
# to 1, for finding the intercept. The second column is the time at
# each group, for finding the slope. The remaining columns (if any),
# are 0 for all rows prior to a certain point, then 1 for all
# subsequent rows (i.e. the Heaviside function). The transition from
# 0 to 1 is the location of a cosmic ray hit; the first 1 in a column
# corresponds to the value in cr_flagged_2d being 1.
x = np.zeros((nz, ngroups, 2 + num_cr), dtype=np.float64)
x[:, :, 0] = 1.
x[:, :, 1] = np.arange(ngroups, dtype=np.float64) * group_time + \
frame_time * (M + 1.) / 2.
if num_cr > 0:
sum_crs = cr_flagged_2d.cumsum(axis=0)
for k in range(ngroups):
s = slice(k, ngroups)
for n in range(1, num_cr + 1):
temp = np.where(
np.logical_and(cr_flagged_2d[k] == 1, sum_crs[k] == n))
if len(temp[0]) > 0:
index = (temp[0], s, n + 1)
x[index] = 1
del temp, index
y = np.transpose(ramp_data, (1, 0)).reshape((nz, ngroups, 1))
# cov is an array of nz matrices, each ngroups x ngroups. The
# inverse of each of these matrices is a weight matrix.
# Note that there are two objects that are called the covariance
# matrix: (1) the ngroups x ngroups matrices in cov, and (2) the
# smaller matrix (see near the end of this function) that contains
# the variances and covariances of the fitted parameters.
cov = np.ones((nz, ngroups, ngroups), dtype=np.float64)
# Use the previous fit to the data to populate the covariance matrix,
# for each of the nz pixels. prev_fit_data has shape (ngroups, nz),
# similar to the ramp data, but we want the nz axis to be the first
# (we're constructing an array of nz matrix equations), so transpose
# prev_fit_data.
prev_fit_T = np.transpose(prev_fit_data, (1, 0))
for k in range(ngroups):
# Populate the upper right, row by row.
cov[:, k, k:ngroups] = prev_fit_T[:, k:k + 1]
# Populate the lower left, column by column.
cov[:, k:ngroups, k] = prev_fit_T[:, k:k + 1]
# Propagate errors from input.
cov[:, k, k] += input_var_data[k, :]
# Give saturated pixels very low weight (i.e. high variance).
cov[:, k, k] += saturated_data[k, :]
del prev_fit_T
# I is 2-D, but it can broadcast to 4-D. This is used to add terms to
# the diagonal of the covariance matrix.
I = np.identity(ngroups)
# Divide by sqrt(2) to convert the readnoise from CDS to single readout.
rn3d = readnoise.reshape((nz, 1, 1)) * SINGLE_READOUT_RN_FACTOR
cov += (I * (rn3d**2 / M))
# prev_slope_data must be non-negative.
flags = prev_slope_data < 0.
prev_slope_data[flags] = 1.
if use_extra_terms:
# Include two dummy axes to allow broadcasting with cov.
slope3d = prev_slope_data.reshape((nz, 1, 1))
# diagonal:
if gain is not None:
g3d = gain.reshape((nz, 1, 1))
else:
g3d = 1.
delta_diag = I * (slope3d * frame_time * (M - 1.) * (M - 2.) /
(3. * M) + (g3d * M)**2 / 12.)
cov += delta_diag
del delta_diag
# This is the solution: (xT @ weight @ x)^-1 @ [xT @ weight @ y]
# where @ means matrix multiplication.
# shape of xT is (nz, 2 + num_cr, ngroups)
xT = np.transpose(x, (0, 2, 1))
# shape of `weight` is (nz, ngroups, ngroups)
I = I.reshape((1, ngroups, ngroups))
weight = la.solve(cov, I) # inverse of cov
del I
# temp1 = xT @ weight
# shape of temp1 is (nz, 2 + num_cr, ngroups)
temp1 = np.einsum('...ij,...jk->...ik', xT, weight)
# temp_var = xT @ weight @ x
# shape of temp_var is (nz, 2 + num_cr, 2 + num_cr)
temp_var = np.einsum('...ij,...jk->...ik', temp1, x)
# `covar` is an array of nz covariance matrices.
# covar = (xT @ weight @ x)^-1
# shape of covar is (nz, 2 + num_cr, 2 + num_cr)
I_2 = np.eye(2 + num_cr).reshape((1, 2 + num_cr, 2 + num_cr))
try:
covar = la.solve(temp_var, I_2) # inverse of temp_var
except la.LinAlgError as msg:
for z in range(nz):
try:
dummy = la.solve(temp_var[z], I_2)
except la.LinAlgError as msg2:
log.warn("singular matrix, z = %d" % z)
raise la.LinAlgError(msg2)
del I_2
# [xT @ weight @ y]
# shape of temp2 is (nz, 2 + num_cr, 1)
temp2 = np.einsum('...ij,...jk->...ik', temp1, y)
# shape of result is (nz, 2 + num_cr, 1)
result = np.einsum('...ij,...jk->...ik', covar, temp2)
r_shape = result.shape
result2d = result.reshape((r_shape[0], r_shape[1]))
del result
# shape of both result2d and variances is (nz, 2 + num_cr)
variances = covar.diagonal(axis1=1, axis2=2).copy()
return (result2d, variances)
def f_logistic(X, params):
""" The logistic function """
return 1/(1+np.exp(-np.einsum('ij,j->i', X, params)))
def post(self):
cd('T300K')
try:
df = pd.read_csv(
"BTE.kappa_scalar",
sep=r"[ \t]+",
header=None,
names=['step', 'kappa'],
engine='python')
ks = np.array(df['kappa'])
plot(
(np.array(df['step']), 'Iteration Step'),
(ks, 'Thermal Conductivity (W/mK)'),
'kappa_scalar.png',
grid=True,
linewidth=2)
except Exception as e:
print(e)
try:
df = pd.read_csv(
"BTE.cumulative_kappa_scalar",
sep=r"[ \t]+",
header=None,
names=['l', 'kappa'],
engine='python')
ks = np.array(df['kappa'])
plot(
(np.array(df['l']),
'Cutoff Mean Free Path for Phonons (Angstrom)'),
(ks, 'Thermal Conductivity (W/mK)'),
'cumulative_kappa_scalar.png',
grid=True,
linewidth=2,
logx=True)
except Exception as e:
print(e)
try:
omega = np.loadtxt('../BTE.omega') / (2.0 * np.pi)
kappa = np.loadtxt('BTE.kappa')[-1, 1:]
kappa = np.einsum('jji', kappa.reshape([3, 3, -1])) / 3.0
plot(
(np.arange(len(omega[0])), 'Band'),
(kappa, 'Thermal Conductivity (W/mK)'),
'kappa_band.png',
grid=True,
linewidth=2)
plot(
(np.arange(len(omega[0])), 'Band'),
(kappa.cumsum(), 'Thermal Conductivity (W/mK)'),
'cumulative_kappa_band.png',
grid=True,
linewidth=2)
except Exception as e:
print(e)
try:
kappa = np.loadtxt('BTE.cumulative_kappaVsOmega_tensor')
with fig("atc_freq.png"):
pl.plot(kappa[:, 0], kappa[:, 1], label="${\kappa_{xx}}$")
pl.plot(kappa[:, 0], kappa[:, 5], label="${\kappa_{xx}}$")
pl.plot(kappa[:, 0], kappa[:, 9], label="${\kappa_{xx}}$")
pl.xlabel("Frequency (THz)")
pl.ylabel("Cumulative Thermal Conductivity(W/mK)")
with fig("tc_freq.png"):
pl.plot(
kappa[:, 0],
np.gradient(kappa[:, 1]),
label="${\kappa_{xx}}$")
pl.plot(
kappa[:, 0],
np.gradient(kappa[:, 5]),
label="${\kappa_{xx}}$")
pl.plot(
kappa[:, 0],
np.gradient(kappa[:, 9]),
label="${\kappa_{xx}}$")
pl.xlabel("Frequency (THz)")
pl.ylabel("Cumulative Thermal Conductivity(W/mK)")
except Exception as e:
print(e)
try:
g = np.loadtxt('../BTE.gruneisen')
y = (g.flatten(), 'Gruneisen')
plot(
(omega.flatten(), 'Frequency (THz)'),
y,
'gruneisen_freq.png',
grid=True,
scatter=True)
with fig('gruneisen_freq.png'):
pl.scatter(
omega.flatten(), g.flatten(), marker='.', color='r', s=50)
pl.xlabel('Frequency (THz)')
pl.ylabel('Gruneisen Coeffecient')
# pl.grid(True)
pl.xlim([0, omega.max()])
pl.ylim([-10, 5])
# pl.tick_params(axis='both', which='major', labelsize=14)
to_txt(['freq', 'gruneisen'],
np.c_[omega.flatten(), g.flatten()], 'gruneisen_freq.txt')
g = np.loadtxt('../BTE.P3')
with fig('p3_freq.png'):
pl.scatter(
omega.flatten(),
g.flatten() * 1e6,
marker='.',
color='r',
s=50)
pl.xlabel('Frequency (THz)')
pl.ylabel('P3 $(\\times 10^{-6})$')
# pl.grid(True)
pl.xlim([0, omega.max()])
pl.ylim([0, g.max() * 1e6])
to_txt(['freq', 'p3'],
np.c_[omega.flatten(), g.flatten()], 'p3_freq.txt')
except Exception as e:
print(e)
self.draw_gv()
self.draw_branch_scatter()
self.draw_tau()
cd('..')
def transform(times, accels, start_times, start_time_indices, alphas,
sensors_pos, lin_resp):
'''Takes time series data as an input and generates a signal value based on
entry and exit 4-vectors on a sphere extended in time. Returns the signal
value and 4-vectors. Refer to Qin's note for a much more detailed
explanation.
accels is a list of accelerations.
sensor_dict is a list of sensor positions, in the same order as accels.
'''
'''The output of this conjoined integral transform function is the S value, not the SNR.
And this value cannot be transformed into the SNR trivially. This is because of how the
function is set up and how the expected signal from the sensor is generated. Here, for
every sensor, an expected signal is generated, analyzed with the data, then discarded.
The result of each run is added continously and the final value is the S, but each
expected signal has been lost and cannot be used for the SNR calculation (see the
separated integral transform for comparison).'''
S = []
S_norm = []
alpha0_x = []
alpha0_y = []
alpha0_z = []
alpha0_t = []
alpha1_x = []
alpha1_y = []
alpha1_z = []
alpha1_t = []
steps = []
adc_timestep_size = times[1] - times[0]
response_length = len(lin_resp)
for i, start_time in enumerate(tqdm(start_times)):
for alpha_index in range(alphas.shape[0]):
alpha_pair = alphas[alpha_index, :]
start_index = start_time_indices[i]
dir_vector = np.array([
alpha_pair[4] - alpha_pair[0],
alpha_pair[5] - alpha_pair[1],
alpha_pair[6] - alpha_pair[2],
])
initial_pos = np.array(
[alpha_pair[0], alpha_pair[1], alpha_pair[2]])
dir_vector_step = dir_vector / (alpha_pair[7] -
alpha_pair[3]) * adc_timestep_size
n_steps = min(
int(
np.ceil(
(alpha_pair[7] - alpha_pair[3]) / adc_timestep_size)),
len(times[times > start_time]))
particle_pos_arr = np.array(
[initial_pos + j * dir_vector_step for j in range(n_steps)])
track_times = np.array(
[start_time + j * adc_timestep_size for j in range(n_steps)])
S_this_track = 0
if n_steps > 0:
for sens_num, sensor_pos in enumerate(sensors_pos):
vector_delta = np.zeros((n_steps, 4))
for j in range(n_steps):
vector_delta[j, 0] = (particle_pos_arr[j][0] -
sensor_pos[0])
vector_delta[j, 1] = (particle_pos_arr[j][1] -
sensor_pos[1])
vector_delta[j, 2] = (particle_pos_arr[j][2] -
sensor_pos[2])
##vector_delta[j, 3] = (track_times[j] - #step_value) # Something like this can be used for analysis of response time
expected_signal_from_sensor = signal_function(
vector_delta, lin_resp, adc_timestep_size)
signal_from_sensor = accels[
sens_num][start_index:start_index +
expected_signal_from_sensor.shape[0]]
S_this_track += np.einsum('ij,ij->',
expected_signal_from_sensor,
signal_from_sensor)
if np.any(np.isnan(S_this_track)):
import pdb
pdb.set_trace()
S.append(S_this_track)
S_norm.append(S_this_track / n_steps)
else:
S.append(0)
S_norm.append(0)
alpha0_x.append(alpha_pair[0])
alpha0_y.append(alpha_pair[1])
alpha0_z.append(alpha_pair[2])
alpha0_t.append(alpha_pair[3] + start_time)
alpha1_x.append(alpha_pair[4])
alpha1_y.append(alpha_pair[5])
alpha1_z.append(alpha_pair[6])
alpha1_t.append(alpha_pair[7] + start_time)
steps.append(n_steps)
structured_array = np.zeros(len(S),
dtype=[
('S', 'f8'),
('S_norm', 'f8'),
('alpha0_x', 'f8'),
('alpha0_y', 'f8'),
('alpha0_z', 'f8'),
('alpha0_t', 'f8'),
('alpha1_x', 'f8'),
('alpha1_y', 'f8'),
('alpha1_z', 'f8'),
('alpha1_t', 'f8'),
('steps', 'i4'),
])
structured_array['S'] = S
structured_array['S_norm'] = S_norm
structured_array['alpha0_x'] = alpha0_x
structured_array['alpha0_y'] = alpha0_y
structured_array['alpha0_z'] = alpha0_z
structured_array['alpha0_t'] = alpha0_t
structured_array['alpha1_x'] = alpha1_x
structured_array['alpha1_y'] = alpha1_y
structured_array['alpha1_z'] = alpha1_z
structured_array['alpha1_t'] = alpha1_t
structured_array['steps'] = steps
return structured_array
def get_jk(mydf, dm, hermi=1, kpt=numpy.zeros(3),
kpt_band=None, with_j=True, with_k=True, exxdiv=None):
'''JK for given k-point'''
from pyscf.pbc.df.df_jk import _ewald_exxdiv_for_G0
vj = vk = None
if kpt_band is not None and abs(kpt-kpt_band).sum() > 1e-9:
kpt = numpy.reshape(kpt, (1,3))
if with_k:
vk = get_k_kpts(mydf, dm, hermi, kpt, kpt_band, exxdiv)
if with_j:
vj = get_j_kpts(mydf, dm, hermi, kpt, kpt_band)
return vj, vk
cell = mydf.cell
log = logger.Logger(mydf.stdout, mydf.verbose)
t1 = (time.clock(), time.time())
dm = numpy.asarray(dm, order='C')
dms = _format_dms(dm, [kpt])
nset, _, nao = dms.shape[:3]
dms = dms.reshape(nset,nao,nao)
j_real = gamma_point(kpt)
k_real = gamma_point(kpt) and not numpy.iscomplexobj(dms)
kptii = numpy.asarray((kpt,kpt))
kpt_allow = numpy.zeros(3)
if with_j:
vjcoulG = mydf.weighted_coulG(kpt_allow, False, mydf.gs)
vjR = numpy.zeros((nset,nao,nao))
vjI = numpy.zeros((nset,nao,nao))
if with_k:
mydf.exxdiv = exxdiv
vkcoulG = mydf.weighted_coulG(kpt_allow, True, mydf.gs)
vkR = numpy.zeros((nset,nao,nao))
vkI = numpy.zeros((nset,nao,nao))
dmsR = numpy.asarray(dms.real.reshape(nset,nao,nao), order='C')
dmsI = numpy.asarray(dms.imag.reshape(nset,nao,nao), order='C')
mem_now = lib.current_memory()[0]
max_memory = max(2000, (mydf.max_memory - mem_now)) * .8
log.debug1('max_memory = %d MB (%d in use)', max_memory, mem_now)
t2 = t1
# rho_rs(-G+k_rs) is computed as conj(rho_{rs^*}(G-k_rs))
# == conj(transpose(rho_sr(G+k_sr), (0,2,1)))
blksize = max(int(max_memory*.25e6/16/nao**2), 16)
bufR = numpy.empty(blksize*nao**2)
bufI = numpy.empty(blksize*nao**2)
for pqkR, pqkI, p0, p1 in mydf.pw_loop(mydf.gs, kptii, max_memory=max_memory):
t2 = log.timer_debug1('%d:%d ft_aopair'%(p0,p1), *t2)
pqkR = pqkR.reshape(nao,nao,-1)
pqkI = pqkI.reshape(nao,nao,-1)
if with_j:
#:v4 = numpy.einsum('ijL,lkL->ijkl', pqk, pqk.conj())
#:vj += numpy.einsum('ijkl,lk->ij', v4, dm)
for i in range(nset):
rhoR = numpy.einsum('pq,pqk->k', dmsR[i], pqkR)
rhoR+= numpy.einsum('pq,pqk->k', dmsI[i], pqkI)
rhoI = numpy.einsum('pq,pqk->k', dmsI[i], pqkR)
rhoI-= numpy.einsum('pq,pqk->k', dmsR[i], pqkI)
rhoR *= vjcoulG[p0:p1]
rhoI *= vjcoulG[p0:p1]
vjR[i] += numpy.einsum('pqk,k->pq', pqkR, rhoR)
vjR[i] -= numpy.einsum('pqk,k->pq', pqkI, rhoI)
if not j_real:
vjI[i] += numpy.einsum('pqk,k->pq', pqkR, rhoI)
vjI[i] += numpy.einsum('pqk,k->pq', pqkI, rhoR)
#t2 = log.timer_debug1(' with_j', *t2)
if with_k:
coulG = numpy.sqrt(vkcoulG[p0:p1])
pqkR *= coulG
pqkI *= coulG
#:v4 = numpy.einsum('ijL,lkL->ijkl', pqk, pqk.conj())
#:vk += numpy.einsum('ijkl,jk->il', v4, dm)
pLqR = lib.transpose(pqkR, axes=(0,2,1), out=bufR).reshape(-1,nao)
pLqI = lib.transpose(pqkI, axes=(0,2,1), out=bufI).reshape(-1,nao)
iLkR = numpy.ndarray((nao*(p1-p0),nao), buffer=pqkR)
iLkI = numpy.ndarray((nao*(p1-p0),nao), buffer=pqkI)
for i in range(nset):
if k_real:
lib.dot(pLqR, dmsR[i], 1, iLkR)
lib.dot(pLqI, dmsR[i], 1, iLkI)
lib.dot(iLkR.reshape(nao,-1), pLqR.reshape(nao,-1).T, 1, vkR[i], 1)
lib.dot(iLkI.reshape(nao,-1), pLqI.reshape(nao,-1).T, 1, vkR[i], 1)
else:
zdotNN(pLqR, pLqI, dmsR[i], dmsI[i], 1, iLkR, iLkI)
zdotNC(iLkR.reshape(nao,-1), iLkI.reshape(nao,-1),
pLqR.reshape(nao,-1).T, pLqI.reshape(nao,-1).T,
1, vkR[i], vkI[i])
#t2 = log.timer_debug1(' with_k', *t2)
pqkR = pqkI = coulG = pLqR = pLqI = iLkR = iLkI = None
#t2 = log.timer_debug1('%d:%d'%(p0,p1), *t2)
bufR = bufI = None
t1 = log.timer_debug1('aft_jk.get_jk', *t1)
if with_j:
if j_real:
vj = vjR
else:
vj = vjR + vjI * 1j
vj = vj.reshape(dm.shape)
if with_k:
if k_real:
vk = vkR
else:
vk = vkR + vkI * 1j
if cell.dimension != 3 and exxdiv is not None:
assert(exxdiv.lower() == 'ewald')
_ewald_exxdiv_for_G0(cell, kpt, dms, vk)
vk = vk.reshape(dm.shape)
return vj, vk
def postold(self):
try:
df = pd.read_csv(
"BTE.kappa_scalar",
sep=r"[ \t]+",
header=None,
names=['step', 'kappa'],
engine='python')
ks = np.array(df['kappa'])
plot(
(np.array(df['step']), 'Iteration Step'),
(ks, 'Thermal Conductivity (W/mK)'),
'kappa_scalar.png',
grid=True,
linewidth=2)
except Exception as e:
print(e)
try:
df = pd.read_csv(
"BTE.cumulative_kappa_scalar",
sep=r"[ \t]+",
header=None,
names=['l', 'kappa'],
engine='python')
ks = np.array(df['kappa'])
plot(
(np.array(df['l']),
'Cutoff Mean Free Path for Phonons (Angstrom)'),
(ks, 'Thermal Conductivity (W/mK)'),
'cumulative_kappa_scalar.png',
grid=True,
linewidth=2,
logx=True)
except Exception as e:
print(e)
try:
omega = np.loadtxt('BTE.omega') / (2.0 * np.pi)
kappa = np.loadtxt('BTE.kappa')[-1, 1:]
kappa = np.einsum('jji', kappa.reshape([3, 3, -1])) / 3.0
plot(
(np.arange(len(omega[0])), 'Band'),
(kappa, 'Thermal Conductivity (W/mK)'),
'kappa_band.png',
grid=True,
linewidth=2)
plot(
(np.arange(len(omega[0])), 'Band'),
(kappa.cumsum(), 'Thermal Conductivity (W/mK)'),
'cumulative_kappa_band.png',
grid=True,
linewidth=2)
except Exception as e:
print(e)
try:
w = np.loadtxt('BTE.w_final')
w = np.abs(w)
w[omega < omega.flatten().max() * 0.005] = float('nan')
plot(
(omega.flatten(), 'Frequency (THz)'), (w.flatten(),
'Scatter Rate (THz)'),
'scatter_freq.png',
grid=True,
scatter=True,
logy=True)
tao = 1.0 / w + 1e-6
with fig('tao_freq.png'):
pl.semilogy(
omega.flatten(),
tao.flatten(),
linestyle='.',
marker='.',
color='r',
markersize=5)
pl.xlabel('Frequency (THz)')
pl.ylabel('Relaxation Time (ps)')
pl.grid(True)
pl.xlim([0, omega.max()])
# pl.ylim([0,tao.flatten().max()])
to_txt(['freq', 'tao'],
np.c_[omega.flatten(), tao.flatten()], 'tao_freq.txt')
except Exception as e:
print(e)
"""
if not exists('relaxtime'):mkdir('relaxtime')
cd('relaxtime')
for i,om in enumerate(omega[:6]):
print "q : ",i
plot((om,'Frequency (THz)'),(tao[i],'Relaxation Time (ps)'),
'tao_freq_q%d.png'%i,grid=True,scatter=True,logx=True,logy=True)
cd('..')
"""
try:
v = np.loadtxt(open('BTE.v'))
n, m = v.shape
v = v.reshape([n, 3, m / 3])
v = np.linalg.norm(v, axis=1)
y = (v.flatten(), 'Group Velocity (nm/ps)')
plot(
(omega.flatten(), 'Frequency (THz)'),
y,
'v_freq.png',
grid=True,
scatter=True)
to_txt(['freq', 'vg'],
np.c_[omega.flatten(), v.flatten()], 'v_freq.txt')
except Exception as e:
print(e)
try:
l = v * tao
l[l < 1e-6] = None
plot(
(omega.flatten(), 'Frequency (THz)'), (l.flatten(),
'Mean Free Path (nm)'),
'lamda_freq.png',
grid=True,
scatter=True,
logy=True,
logx=True,
xmin=0)
to_txt(['freq', 'mfp'],
np.c_[omega.flatten(), l.flatten()], 'lamda_freq.txt')
except Exception as e:
print(e)
try:
q = np.loadtxt(open('BTE.qpoints'))
qnorm = np.linalg.norm(q[:, -3:], axis=1)
data = []
n, m = w.shape
for i in range(m):
data.append([qnorm, w[:, i], 'b'])
series(
xlabel='|q| (1/nm)',
ylabel='Scatter Rate (THz)',
datas=data,
filename='branchscatter.png',
scatter=True,
legend=False,
logx=True,
logy=True)
except Exception as e:
print(e)
def partial_hess_elec(hessobj,
mo_energy=None,
mo_coeff=None,
mo_occ=None,
atmlst=None,
max_memory=4000,
verbose=None):
'''Partial derivative
'''
log = logger.new_logger(hessobj, verbose)
time0 = t1 = (time.clock(), time.time())
mol = hessobj.mol
mf = hessobj.base
if mo_energy is None: mo_energy = mf.mo_energy
if mo_occ is None: mo_occ = mf.mo_occ
if mo_coeff is None: mo_coeff = mf.mo_coeff
if atmlst is None: atmlst = range(mol.natm)
nao, nmo = mo_coeff.shape
mocc = mo_coeff[:, mo_occ > 0]
nocc = mocc.shape[1]
dm0 = numpy.dot(mocc, mocc.T) * 2
# Energy weighted density matrix
dme0 = numpy.einsum('pi,qi,i->pq', mocc, mocc, mo_energy[mo_occ > 0]) * 2
hcore_deriv = hessobj.hcore_generator(mol)
s1aa, s1ab, s1a = get_ovlp(mol)
vj1, vk1 = _get_jk(
mol,
'int2e_ipip1',
9,
's2kl',
[
'lk->s1ij',
dm0, # vj1
'jk->s1il',
dm0
]) # vk1
vhf_diag = vj1 - vk1 * .5
vhf_diag = vhf_diag.reshape(3, 3, nao, nao)
vj1 = vk1 = None
t1 = log.timer_debug1('contracting int2e_ipip1', *t1)
aoslices = mol.aoslice_by_atom()
de2 = numpy.zeros((mol.natm, mol.natm, 3, 3)) # (A,B,dR_A,dR_B)
for i0, ia in enumerate(atmlst):
shl0, shl1, p0, p1 = aoslices[ia]
shls_slice = (shl0, shl1) + (0, mol.nbas) * 3
vj1, vk1, vk2 = _get_jk(
mol,
'int2e_ip1ip2',
9,
's1',
[
'ji->s1kl',
dm0[:, p0:p1], # vj1
'li->s1kj',
dm0[:, p0:p1], # vk1
'lj->s1ki',
dm0
], # vk2
shls_slice=shls_slice)
vhf = vj1 * 2 - vk1 * .5
vhf[:, :, p0:p1] -= vk2 * .5
t1 = log.timer_debug1('contracting int2e_ip1ip2 for atom %d' % ia, *t1)
vj1, vk1 = _get_jk(
mol,
'int2e_ipvip1',
9,
's2kl',
[
'lk->s1ij',
dm0, # vj1
'li->s1kj',
dm0[:, p0:p1]
], # vk1
shls_slice=shls_slice)
vhf[:, :, p0:p1] += vj1.transpose(0, 2, 1)
vhf -= vk1.transpose(0, 2, 1) * .5
vhf = vhf.reshape(3, 3, nao, nao)
t1 = log.timer_debug1('contracting int2e_ipvip1 for atom %d' % ia, *t1)
s1ao = numpy.zeros((3, nao, nao))
s1ao[:, p0:p1] += s1a[:, p0:p1]
s1ao[:, :, p0:p1] += s1a[:, p0:p1].transpose(0, 2, 1)
s1oo = numpy.einsum('xpq,pi,qj->xij', s1ao, mocc, mocc)
de2[i0, i0] += numpy.einsum('xypq,pq->xy', vhf_diag[:, :, p0:p1],
dm0[p0:p1]) * 2
de2[i0, i0] -= numpy.einsum('xypq,pq->xy', s1aa[:, :, p0:p1],
dme0[p0:p1]) * 2
for j0, ja in enumerate(atmlst[:i0 + 1]):
q0, q1 = aoslices[ja][2:]
# *2 for +c.c.
de2[i0, j0] += numpy.einsum('xypq,pq->xy', vhf[:, :, q0:q1],
dm0[q0:q1]) * 2
de2[i0,
j0] -= numpy.einsum('xypq,pq->xy', s1ab[:, :, p0:p1, q0:q1],
dme0[p0:p1, q0:q1]) * 2
h1ao = hcore_deriv(ia, ja)
de2[i0, j0] += numpy.einsum('xypq,pq->xy', h1ao, dm0)
for j0 in range(i0):
de2[j0, i0] = de2[i0, j0].T
log.timer('RHF partial hessian', *time0)
return de2
def MK(u,k,e,r,mu,ReTau,mesh,compressibleCorrection):
import numpy as np
from solveEqn import solveEqn
n = mesh.nPoints
d = np.minimum(mesh.y, mesh.y[-1]-mesh.y)
if compressibleCorrection == 1:
yplus = d*np.sqrt(r/r[0])/(mu/mu[0])*ReTau
else:
yplus = d*ReTau
# Model constants
cmu = 0.09
sigk = 1.4
sige = 1.3
Ce1 = 1.4
Ce2 = 1.8
# Model functions
ReTurb = r*np.power(k, 2)/(mu*e)
f2 = (1-2/9*np.exp(-np.power(ReTurb/6, 2)))*np.power(1-np.exp(-yplus/5), 2)
fmue = (1-np.exp(-yplus/70))*(1.0+3.45/np.power(ReTurb, 0.5))
fmue[0] = fmue[-1] = 0.0
# eddy viscosity
mut = cmu*fmue*r/e*np.power(k,2)
mut[1:-1] = np.minimum(np.maximum(mut[1:-1],1.0e-10),100.0)
# Turbulent production: Pk = mut*dudy^2
Pk = mut*np.power(mesh.ddy@u, 2)
# ---------------------------------------------------------------------
# e-equation
# effective viscosity
if compressibleCorrection == 1:
mueff = (mu + mut/sige)/np.sqrt(r)
fs = np.power(r, 1.5)
fd = 1/r
else:
mueff = mu + mut/sige
fs = fd = np.ones(n)
# diffusion matrix: mueff*d2()/dy2 + dmueff/dy d()/dy
A = np.einsum('i,ij->ij', mueff*fd, mesh.d2dy2) \
+ np.einsum('i,ij->ij', (mesh.ddy@mueff)*fd, mesh.ddy)
# Left-hand-side, implicitly treated source term
np.fill_diagonal(A, A.diagonal() - Ce2*f2*r*e/k/fs)
# Right-hand-side
b = -e[1:-1]/k[1:-1]*Ce1*Pk[1:-1]
# Wall boundary conditions
e[0 ] = mu[ 0]/r[ 0]*k[ 1]/np.power(d[ 1], 2)
e[-1] = mu[-1]/r[-1]*k[-2]/np.power(d[-2], 2)
# Solve eps equation
e = solveEqn(e*fs, A, b, 0.8)/fs
e[1:-1] = np.maximum(e[1:-1], 1.e-12)
# ---------------------------------------------------------------------
# k-equation
# effective viscosity
if compressibleCorrection == 1:
mueff = (mu + mut/sigk)/np.sqrt(r)
fs = r
fd = 1/np.sqrt(r)
else:
mueff = mu + mut/sigk
fs = fd = np.ones(n)
# diffusion matrix: mueff*d2()/dy2 + dmueff/dy d()/dy
A = np.einsum('i,ij->ij', mueff*fd, mesh.d2dy2) \
+ np.einsum('i,ij->ij', (mesh.ddy@mueff)*fd, mesh.ddy)
# implicitly treated source term
np.fill_diagonal(A, A.diagonal() - r*e/k/fs)
# Right-hand-side
b = -Pk[1:-1]
# Wall boundary conditions
k[0] = k[-1] = 0.0
# Solve TKE
k = solveEqn(k*fs, A, b, 0.7)/fs
k[1:-1] = np.maximum(k[1:-1], 1.e-12)
return mut,k,e
def solve_mo1(mf,
mo_energy,
mo_coeff,
mo_occ,
h1ao_or_chkfile,
fx=None,
atmlst=None,
max_memory=4000,
verbose=None):
mol = mf.mol
if atmlst is None: atmlst = range(mol.natm)
nao, nmo = mo_coeff.shape
mocc = mo_coeff[:, mo_occ > 0]
nocc = mocc.shape[1]
if fx is None:
fx = gen_vind(mf, mo_coeff, mo_occ)
s1a = -mol.intor('int1e_ipovlp', comp=3)
def _ao2mo(mat):
return numpy.asarray(
[reduce(numpy.dot, (mo_coeff.T, x, mocc)) for x in mat])
mem_now = lib.current_memory()[0]
max_memory = max(2000, max_memory * .9 - mem_now)
blksize = max(2, int(max_memory * 1e6 / 8 / (nmo * nocc * 3 * 6)))
mo1s = [None] * mol.natm
e1s = [None] * mol.natm
aoslices = mol.aoslice_by_atom()
for ia0, ia1 in lib.prange(0, len(atmlst), blksize):
s1vo = []
h1vo = []
for i0 in range(ia0, ia1):
ia = atmlst[i0]
shl0, shl1, p0, p1 = aoslices[ia]
s1ao = numpy.zeros((3, nao, nao))
s1ao[:, p0:p1] += s1a[:, p0:p1]
s1ao[:, :, p0:p1] += s1a[:, p0:p1].transpose(0, 2, 1)
s1vo.append(_ao2mo(s1ao))
if isinstance(h1ao_or_chkfile, str):
key = 'scf_f1ao/%d' % ia
h1ao = lib.chkfile.load(h1ao_or_chkfile, key)
else:
h1ao = h1ao_or_chkfile[ia]
h1vo.append(_ao2mo(h1ao))
h1vo = numpy.vstack(h1vo)
s1vo = numpy.vstack(s1vo)
mo1, e1 = cphf.solve(fx, mo_energy, mo_occ, h1vo, s1vo)
mo1 = numpy.einsum('pq,xqi->xpi', mo_coeff,
mo1).reshape(-1, 3, nao, nocc)
e1 = e1.reshape(-1, 3, nocc, nocc)
for k in range(ia1 - ia0):
ia = atmlst[k + ia0]
if isinstance(h1ao_or_chkfile, str):
key = 'scf_mo1/%d' % ia
lib.chkfile.save(h1ao_or_chkfile, key, mo1[k])
else:
mo1s[ia] = mo1[k]
e1s[ia] = e1[k].reshape(3, nocc, nocc)
mo1 = e1 = None
if isinstance(h1ao_or_chkfile, str):
return h1ao_or_chkfile, e1s
else:
return mo1s, e1s
if __name__ == '__main__':
from pyscf.pbc import gto as pgto
from pyscf.pbc import scf as pscf
from pyscf.pbc.df import aft
L = 5.
n = 5
cell = pgto.Cell()
cell.a = numpy.diag([L,L,L])
cell.gs = numpy.array([n,n,n])
cell.atom = '''He 3. 2. 3.
He 1. 1. 1.'''
#cell.basis = {'He': [[0, (1.0, 1.0)]]}
#cell.basis = '631g'
#cell.basis = {'He': [[0, (2.4, 1)], [1, (1.1, 1)]]}
cell.basis = 'ccpvdz'
cell.verbose = 0
cell.build(0,0)
cell.verbose = 5
df = aft.AFTDF(cell)
df.gs = (15,)*3
dm = pscf.RHF(cell).get_init_guess()
vj, vk = df.get_jk(dm)
print(numpy.einsum('ij,ji->', df.get_nuc(), dm), 'ref=-10.577490961074622')
print(numpy.einsum('ij,ji->', vj, dm), 'ref=5.3766911667862516')
print(numpy.einsum('ij,ji->', vk, dm), 'ref=8.2255177602309022')
def make_soc2e(gobj, dm0):
dma, dmb = dm0
nao = dma.shape[0]
# FIXME: see JPC, 101, 3388 Eq (11c), why?
g_so = (lib.param.G_ELECTRON - 1) * 2
# hso2e is the imaginary part of SSO
hso2e = mol.intor('int2e_p1vxp1', 3).reshape(3, nao, nao, nao, nao)
vj = numpy.zeros((2, 3, nao, nao))
vk = numpy.zeros((2, 3, nao, nao))
if gobj.with_sso:
vj[:] += g_so / 2 * numpy.einsum('yijkl,ji->ykl', hso2e, dma - dmb)
vj[0] += g_so / 2 * numpy.einsum('yijkl,lk->yij', hso2e, dma + dmb)
vj[1] -= g_so / 2 * numpy.einsum('yijkl,lk->yij', hso2e, dma + dmb)
vk[0] += g_so / 2 * numpy.einsum('yijkl,jk->yil', hso2e, dma)
vk[1] -= g_so / 2 * numpy.einsum('yijkl,jk->yil', hso2e, dmb)
vk[0] += g_so / 2 * numpy.einsum('yijkl,li->ykj', hso2e, dma)
vk[0] -= g_so / 2 * numpy.einsum('yijkl,li->ykj', hso2e, dmb)
if gobj.with_soo:
vj[0] += 2 * numpy.einsum('yijkl,ji->ykl', hso2e, dma + dmb)
vj[1] -= 2 * numpy.einsum('yijkl,ji->ykl', hso2e, dma + dmb)
vj[:] += 2 * numpy.einsum('yijkl,lk->yij', hso2e, dma - dmb)
vk[0] += 2 * numpy.einsum('yijkl,jk->yil', hso2e, dma)
vk[1] -= 2 * numpy.einsum('yijkl,jk->yil', hso2e, dmb)
vk[0] += 2 * numpy.einsum('yijkl,li->ykj', hso2e, dma)
vk[1] -= 2 * numpy.einsum('yijkl,li->ykj', hso2e, dmb)
hso2e = vj - vk
return hso2e
def gen_hop(hobj, mo_energy=None, mo_coeff=None, mo_occ=None, verbose=None):
log = logger.new_logger(hobj, verbose)
mol = hobj.mol
mf = hobj.base
if mo_energy is None: mo_energy = mf.mo_energy
if mo_occ is None: mo_occ = mf.mo_occ
if mo_coeff is None: mo_coeff = mf.mo_coeff
natm = mol.natm
nao, nmo = mo_coeff.shape
mocc = mo_coeff[:, mo_occ > 0]
nocc = mocc.shape[1]
atmlst = range(natm)
max_memory = max(2000, hobj.max_memory - lib.current_memory()[0])
de2 = hobj.partial_hess_elec(mo_energy, mo_coeff, mo_occ, atmlst,
max_memory, log)
de2 += hobj.hess_nuc()
# Compute H1 integrals and store in hobj.chkfile
hobj.make_h1(mo_coeff, mo_occ, hobj.chkfile, atmlst, log)
aoslices = mol.aoslice_by_atom()
s1a = -mol.intor('int1e_ipovlp', comp=3)
fvind = gen_vind(mf, mo_coeff, mo_occ)
def h_op(x):
x = x.reshape(natm, 3)
hx = numpy.einsum('abxy,ax->by', de2, x)
h1ao = 0
s1ao = 0
for ia in range(natm):
shl0, shl1, p0, p1 = aoslices[ia]
h1ao_i = lib.chkfile.load(hobj.chkfile, 'scf_f1ao/%d' % ia)
h1ao += numpy.einsum('x,xij->ij', x[ia], h1ao_i)
s1ao_i = numpy.zeros((3, nao, nao))
s1ao_i[:, p0:p1] += s1a[:, p0:p1]
s1ao_i[:, :, p0:p1] += s1a[:, p0:p1].transpose(0, 2, 1)
s1ao += numpy.einsum('x,xij->ij', x[ia], s1ao_i)
s1vo = reduce(numpy.dot, (mo_coeff.T, s1ao, mocc))
h1vo = reduce(numpy.dot, (mo_coeff.T, h1ao, mocc))
mo1, mo_e1 = cphf.solve(fvind, mo_energy, mo_occ, h1vo, s1vo)
mo1 = numpy.dot(mo_coeff, mo1)
mo_e1 = mo_e1.reshape(nocc, nocc)
dm1 = numpy.einsum('pi,qi->pq', mo1, mocc)
dme1 = numpy.einsum('pi,qi,i->pq', mo1, mocc, mo_energy[mo_occ > 0])
dme1 = dme1 + dme1.T + reduce(numpy.dot, (mocc, mo_e1, mocc.T))
for ja in range(natm):
q0, q1 = aoslices[ja][2:]
h1ao = lib.chkfile.load(hobj.chkfile, 'scf_f1ao/%s' % ja)
hx[ja] += numpy.einsum('xpq,pq->x', h1ao, dm1) * 4
hx[ja] -= numpy.einsum('xpq,pq->x', s1a[:, q0:q1], dme1[q0:q1]) * 2
hx[ja] -= numpy.einsum('xpq,qp->x', s1a[:, q0:q1], dme1[:,
q0:q1]) * 2
return hx.ravel()
hdiag = numpy.einsum('aaxx->ax', de2).ravel()
return h_op, hdiag
def _get_e_T_Emnar_2(self, eps_Emab):
# get the tangential strain vector array for each microplane
MPTT_ijr = self._get__MPTT()
return np.einsum('nija,...ij->...na', MPTT_ijr, eps_Emab)
def hess_elec(hessobj,
mo_energy=None,
mo_coeff=None,
mo_occ=None,
mo1=None,
mo_e1=None,
h1ao=None,
atmlst=None,
max_memory=4000,
verbose=None):
log = logger.new_logger(hessobj, verbose)
time0 = t1 = (time.clock(), time.time())
mol = hessobj.mol
mf = hessobj.base
if mo_energy is None: mo_energy = mf.mo_energy
if mo_occ is None: mo_occ = mf.mo_occ
if mo_coeff is None: mo_coeff = mf.mo_coeff
if atmlst is None: atmlst = range(mol.natm)
de2 = hessobj.partial_hess_elec(mo_energy, mo_coeff, mo_occ, atmlst,
max_memory, log)
if h1ao is None:
h1ao = hessobj.make_h1(mo_coeff, mo_occ, hessobj.chkfile, atmlst, log)
t1 = log.timer_debug1('making H1', *time0)
if mo1 is None or mo_e1 is None:
mo1, mo_e1 = hessobj.solve_mo1(mo_energy, mo_coeff, mo_occ, h1ao, None,
atmlst, max_memory, log)
t1 = log.timer_debug1('solving MO1', *t1)
if isinstance(h1ao, str):
h1ao = lib.chkfile.load(h1ao, 'scf_f1ao')
h1ao = dict([(int(k), h1ao[k]) for k in h1ao])
if isinstance(mo1, str):
mo1 = lib.chkfile.load(mo1, 'scf_mo1')
mo1 = dict([(int(k), mo1[k]) for k in mo1])
nao, nmo = mo_coeff.shape
mocc = mo_coeff[:, mo_occ > 0]
s1a = -mol.intor('int1e_ipovlp', comp=3)
aoslices = mol.aoslice_by_atom()
for i0, ia in enumerate(atmlst):
shl0, shl1, p0, p1 = aoslices[ia]
s1ao = numpy.zeros((3, nao, nao))
s1ao[:, p0:p1] += s1a[:, p0:p1]
s1ao[:, :, p0:p1] += s1a[:, p0:p1].transpose(0, 2, 1)
s1oo = numpy.einsum('xpq,pi,qj->xij', s1ao, mocc, mocc)
for j0 in range(i0 + 1):
ja = atmlst[j0]
q0, q1 = aoslices[ja][2:]
# *2 for double occupancy, *2 for +c.c.
dm1 = numpy.einsum('ypi,qi->ypq', mo1[ja], mocc)
de2[i0, j0] += numpy.einsum('xpq,ypq->xy', h1ao[ia], dm1) * 4
dm1 = numpy.einsum('ypi,qi,i->ypq', mo1[ja], mocc,
mo_energy[mo_occ > 0])
de2[i0, j0] -= numpy.einsum('xpq,ypq->xy', s1ao, dm1) * 4
de2[i0, j0] -= numpy.einsum('xpq,ypq->xy', s1oo, mo_e1[ja]) * 2
for j0 in range(i0):
de2[j0, i0] = de2[i0, j0].T
log.timer('RHF hessian', *time0)
return de2
def _get__MPNN(self):
MPNN_nij = np.einsum('ni,nj->nij', self._MPN, self._MPN)
return MPNN_nij
def test_ft_aoao(self):
#coords = pdft.gen_grid.gen_uniform_grids(cell)
#aoR = pdft.numint.eval_ao(cell, coords)
#ngs, nao = aoR.shape
#ref = numpy.asarray([tools.fft(aoR[:,i].conj()*aoR[:,j], cell.gs)
# for i in range(nao) for j in range(nao)])
#ref = ref.reshape(nao,nao,-1).transpose(2,0,1) * (cell.vol/ngs)
#dat = ft_ao.ft_aopair(cell, cell.Gv, aosym='s1hermi')
#self.assertAlmostEqual(numpy.linalg.norm(ref[:,0,0]-dat[:,0,0]) , 0, 5)
#self.assertAlmostEqual(numpy.linalg.norm(ref[:,1,1]-dat[:,1,1]) , 0.02315483195832373, 4)
#self.assertAlmostEqual(numpy.linalg.norm(ref[:,2:,2:]-dat[:,2:,2:]), 0, 9)
#self.assertAlmostEqual(numpy.linalg.norm(ref[:,0,2:]-dat[:,0,2:]) , 0, 9)
#self.assertAlmostEqual(numpy.linalg.norm(ref[:,2:,0]-dat[:,2:,0]) , 0, 9)
#idx = numpy.tril_indices(nao)
#ref = dat[:,idx[0],idx[1]]
#dat = ft_ao.ft_aopair(cell, cell.Gv, aosym='s2')
#self.assertAlmostEqual(abs(dat-ref).sum(), 0, 9)
coords = pdft.gen_grid.gen_uniform_grids(cell1)
Gv, Gvbase, kws = cell1.get_Gv_weights(cell1.gs)
b = cell1.reciprocal_vectors()
gxyz = lib.cartesian_prod([numpy.arange(len(x)) for x in Gvbase])
dat = ft_ao.ft_aopair(cell1,
cell1.Gv,
aosym='s1',
b=b,
gxyz=gxyz,
Gvbase=Gvbase)
self.assertAlmostEqual(finger(dat),
1.5666516306798806 + 1.953555017583245j, 9)
dat = ft_ao.ft_aopair(cell1,
cell1.Gv,
aosym='s2',
b=b,
gxyz=gxyz,
Gvbase=Gvbase)
self.assertAlmostEqual(finger(dat),
-0.85276967757297917 + 1.0378751267506394j, 9)
dat = ft_ao.ft_aopair(cell1,
cell1.Gv,
aosym='s1hermi',
b=b,
gxyz=gxyz,
Gvbase=Gvbase)
self.assertAlmostEqual(finger(dat),
1.5666516306798806 + 1.953555017583245j, 9)
aoR = pdft.numint.eval_ao(cell1, coords)
ngrids, nao = aoR.shape
aoaoR = numpy.einsum('pi,pj->ijp', aoR, aoR)
ref = tools.fft(aoaoR.reshape(nao * nao, -1), cell1.gs)
ref = ref.reshape(nao, nao, -1).transpose(2, 0,
1) * (cell1.vol / ngrids)
self.assertAlmostEqual(numpy.linalg.norm(ref[:, 0, 0] - dat[:, 0, 0]),
0, 7)
self.assertAlmostEqual(numpy.linalg.norm(ref[:, 1, 1] - dat[:, 1, 1]),
0, 7)
self.assertAlmostEqual(
numpy.linalg.norm(ref[:, 2:, 2:] - dat[:, 2:, 2:]), 0, 7)
self.assertAlmostEqual(
numpy.linalg.norm(ref[:, 0, 2:] - dat[:, 0, 2:]), 0, 7)
self.assertAlmostEqual(
numpy.linalg.norm(ref[:, 2:, 0] - dat[:, 2:, 0]), 0, 7)
idx = numpy.tril_indices(nao)
ref = dat[:, idx[0], idx[1]]
dat = ft_ao.ft_aopair(cell1, cell1.Gv, aosym='s2')
self.assertAlmostEqual(abs(dat - ref).sum(), 0, 9)
def get_tangential_law(self, eps_T_Emna, omega_T_Emn, z_T_Emn,
alpha_T_Emna, eps_T_pi_Emna, sigma_N_Emn):
E_T = self.E * (1.0 - 4 * self.nu) / \
((1.0 + self.nu) * (1.0 - 2 * self.nu))
# thermo forces
sig_pi_trial = E_T * (eps_T_Emna - eps_T_pi_Emna)
Z = self.K_T * z_T_Emn
X = self.gamma_T * alpha_T_Emna
norm_1 = np.sqrt(
np.einsum('...na,...na->...n', (sig_pi_trial - X),
(sig_pi_trial - X)))
Y = 0.5 * E_T * \
np.einsum(
'...na,...na->...n',
(eps_T_Emna - eps_T_pi_Emna),
(eps_T_Emna - eps_T_pi_Emna))
# threshold
f = norm_1 - self.sigma_T_0 - Z + self.m_T * sigma_N_Emn
plas_1 = f > 1e-6
elas_1 = f < 1e-6
delta_lamda = f / \
(E_T / (1.0 - omega_T_Emn) + self.gamma_T + self.K_T) * plas_1
norm_2 = 1.0 * elas_1 + np.sqrt(
np.einsum('...na,...na->...n', (sig_pi_trial - X),
(sig_pi_trial - X))) * plas_1
eps_T_pi_Emna[..., 0] += plas_1 * delta_lamda * \
((sig_pi_trial[..., 0] - X[..., 0]) /
(1.0 - omega_T_Emn)) / norm_2
eps_T_pi_Emna[..., 1] += plas_1 * delta_lamda * \
((sig_pi_trial[..., 1] - X[..., 1]) /
(1.0 - omega_T_Emn)) / norm_2
eps_T_pi_Emna[..., 2] += plas_1 * delta_lamda * \
((sig_pi_trial[..., 2] - X[..., 2]) /
(1.0 - omega_T_Emn)) / norm_2
omega_T_Emn += ((1 - omega_T_Emn) ** self.c_T) * \
(delta_lamda * (Y / self.S_T) ** self.r_T) * \
(self.sigma_T_0 / (self.sigma_T_0 + self.m_T * sigma_N_Emn)) ** self.p_T
omega_T_Emn[...] = np.clip(omega_T_Emn, 0, 0.9999)
alpha_T_Emna[..., 0] += plas_1 * delta_lamda * \
(sig_pi_trial[..., 0] - X[..., 0]) / norm_2
alpha_T_Emna[..., 1] += plas_1 * delta_lamda * \
(sig_pi_trial[..., 1] - X[..., 1]) / norm_2
alpha_T_Emna[..., 2] += plas_1 * delta_lamda * \
(sig_pi_trial[..., 2] - X[..., 2]) / norm_2
z_T_Emn += delta_lamda
sigma_T_Emna = np.einsum('...n,...na->...na', (1 - omega_T_Emn),
E_T * (eps_T_Emna - eps_T_pi_Emna))
Z = self.K_T * z_T_Emn
X = self.gamma_T * alpha_T_Emna
Y = 0.5 * E_T * \
np.einsum(
'...na,...na->...n',
(eps_T_Emna - eps_T_pi_Emna),
(eps_T_Emna - eps_T_pi_Emna))
return sigma_T_Emna, Z, X, Y
def _get_e_N_Emn_2(self, eps_Emab):
# get the normal strain array for each microplane
return np.einsum('nij,...ij->...n', self._MPNN, eps_Emab)
