def TE_two_sites(bonds,vertices,TEO,i,N,d,chi):
"""
-+THETA+-
| |
+-TEO-+
| |
"""
#coarse grain
theta = np.einsum("ij,jsk->isk",np.diag(bonds[(i-1)%N][:]),vertices[i][:,:,:])
theta = np.einsum("isj,jk->isk",theta,np.diag(bonds[i][:]))
theta = np.einsum("isj,jtk->istk",theta,vertices[(i+1)%N][:,:,:])
theta = np.einsum("istj,jk->istk",theta,np.diag(bonds[(i+1)%N][:]))
#multiple operator with wavefunction
theta = np.einsum("istk,stuv->iuvk",theta,TEO)
theta = np.reshape(theta,(chi*d,d*chi))
#svd
X,Y,Z = np.linalg.svd(theta)
bonds[i][0:chi] = Y[0:chi]/np.sqrt(np.sum(Y[0:chi]**2))
X = np.reshape(X[0:chi*d,0:chi],(chi,d,chi))
vertices[i][:,:,:] = np.tensordot(np.diag(bonds[(i-1)%N][:]**(-1)),X,axes=(1,0))
Z= np.reshape(Z[0:chi,0:d*chi],(chi,d,chi))
vertices[(i+1)%N][:,:,:] = np.tensordot(Z,np.diag(bonds[(i+1)%N][:]**(-1)),axes=(2,0))
return theta
def matrix_Heff_optimization(mat_M, mat_W, mat_LR, l):
L = mat_LR[l - 1]
W = mat_W[l]
R = mat_LR[l]
#print '[matrix_Heff_opt] site %i | shapes L/R/W: ' % l, L.shape, R.shape, W.shape
dim_alm1 = L.shape[0]
dim_blm1 = L.shape[1]
dim_bl = W.shape[1]
dim_sigmal = W.shape[3]
dim_al = R.shape[0]
# first tensordot
dim_ind_1 = L.shape[1]
dim_ind_2 = W.shape[0]
if dim_ind_1 != dim_ind_2:
exit('ERROR[matrix_Heff] trying to contract %i-dim and %i-dim indices' % (dim_ind_1, dim_ind_2))
LW = numpy.tensordot(L, W, axes=([1, 0]))
# second tensordot
dim_ind_1 = LW.shape[2]
dim_ind_2 = R.shape[1]
if dim_ind_1 != dim_ind_2:
exit('ERROR[matrix_Heff] trying to contract %i-dim and %i-dim indices' % (dim_ind_1, dim_ind_2))
LWR = numpy.tensordot(LW, R, axes=([2, 1]))
Heff = LWR.reshape(dim_sigmal * dim_alm1 * dim_al, dim_sigmal * dim_alm1 * dim_al)
#Mvec=M.reshape(dim_alm1*dim_sigmal*dim_al,1)
[eigval, eigvec] = scipy.linalg.eigh(Heff, eigvals=(0, 0))
mat_M[l] = eigvec.reshape(dim_alm1, dim_sigmal, dim_al) # modified (tc): mat_M --> mat_M[l]
return eigval[0], mat_M # modified (tc): eigval --> eigval[0]
def project(self,op,v): if(self.conserve): vtemp=v.shallow_copy() vtemp.q_conj=np.array([-1,1]) return npc.tensordot(npc.tensordot(vtemp,op,([0],[0])),v,([1],[0])) else: return np.tensordot(np.tensordot(v,op,([0],[0])),v,([1],[0]))
def backward_cpu(self, inputs, grad_outputs):
x, W = inputs[:2]
b = inputs[2] if len(inputs) == 3 else None
if not type_check.same_types(*inputs):
if b is not None:
raise ValueError('numpy and cupy must not be used together\n'
'type(W): {0}, type(x): {1}, type(b): {2}'
.format(type(W), type(x), type(b)))
else:
raise ValueError('numpy and cupy must not be used together\n'
'type(W): {0}, type(x): {1}'
.format(type(W), type(x)))
gy = grad_outputs[0]
h, w = x.shape[2:]
gW = numpy.tensordot(
gy, self.col, ((0, 2, 3), (0, 4, 5))).astype(W.dtype, copy=False)
if not self.requires_x_grad:
gx = None
else:
gcol = numpy.tensordot(W, gy, (0, 1)).astype(x.dtype, copy=False)
gcol = numpy.rollaxis(gcol, 3)
gx = conv.col2im_cpu(gcol, self.sy, self.sx, self.ph, self.pw,
h, w)
if b is None:
return gx, gW
else:
gb = gy.sum(axis=(0, 2, 3))
return gx, gW, gb
def r_log_spiral(self,phi): """ return distance from center for angle phi of logarithmic spiral Parameters ---------- phi: scalar or np.array with polar angle values Returns ------- r(phi) = rx * exp(b * phi) as np.array Notes ----- see http://en.wikipedia.org/wiki/Logarithmic_spiral """ if np.isscalar(phi): phi = np.array([phi]) ones = np.ones(phi.shape[0]) # self.rx.shape = 8 # phi.shape = p # then result is given as (8,p)-dim array, each row stands for one rx result = np.tensordot(self.rx , np.exp((phi - 3.*pi*ones) / np.tan(pi/2. - self.idisk)),axes = 0) result = np.vstack((result, np.tensordot(self.rx , np.exp((phi - pi*ones) / np.tan(pi/2. - self.idisk)),axes = 0) )) result = np.vstack((result, np.tensordot(self.rx , np.exp((phi + pi*ones) / np.tan(pi/2. - self.idisk)),axes = 0) )) return np.vstack((result, np.tensordot(self.rx , np.exp((phi + 3.*pi*ones) / np.tan(pi/2. - self.idisk)),axes = 0) ))
def contract_two_to_two(self, subscripts, two, out=None, factor=1.0, clear=True):
"""Contract self with a two-index to obtain a two-index.
**Arguments:**
subscripts
Any of ``abcd,bd->ac`` (direct), ``abcd,cb->ad`` (exchange)
two
The input two-index object. (DenseTwoIndex)
**Optional arguments:**
out, factor, clear
See :py:meth:`DenseLinalgFactory.einsum`
"""
check_options('subscripts', subscripts, 'abcd,bd->ac', 'abcd,cb->ad')
if out is None:
out = DenseTwoIndex(self.nbasis)
if clear:
out.clear()
else:
check_type('out', out, DenseTwoIndex)
if subscripts == 'abcd,bd->ac':
tmp = np.tensordot(self._array2, two._array, axes=([(1,2),(1,0)]))
out._array[:] += factor*np.tensordot(self._array, tmp, [0,0])
elif subscripts == 'abcd,cb->ad':
tmp = np.tensordot(self._array2, two._array, axes=([1,1]))
out._array[:] += factor*np.tensordot(self._array, tmp, ([0,2],[0,2]))
return out
def innerProductOBC(mpsA,mpsB):
""" Inner product <A|B> using transfer matrices
where A and B are MPS representations of }A> and }B>
with open boundary conditions (OBC).
"""
# Take adjoint of |A> to get <A|
A = []
for a in mpsA:
A.append(np.conj(a))
B = mpsB
N = len(A) # Number of qubits
d = A[1].shape[1] # d = 2 for qubits
# Construct list of transfer matrices by contracting pairs of
# tensors from A and B.
transfer = []
t = np.tensordot(A[0],B[0],axes=(0,0))
t = np.reshape(t,A[0].shape[1]*B[0].shape[1])
transfer.append(t)
for i in xrange(1,N-1):
t = np.tensordot(A[i],B[i],axes=(1,1))
t = np.transpose(t,axes=(0,2,1,3))
t = np.reshape(t,(A[i].shape[0]*B[i].shape[0],
A[i].shape[2]*B[i].shape[2]))
transfer.append(t)
t = np.tensordot(A[N-1],B[N-1],axes=(1,1))
t = np.reshape(t,A[N-1].shape[0]*B[N-1].shape[0])
transfer.append(t)
# Contract the transfer matrices.
prod = transfer[0]
for i in xrange(1,N-1):
prod = np.dot(prod,transfer[i])
prod = np.dot(prod,transfer[N-1])
return prod
def test_einsum_fixedstridebug(self):
# Issue #4485 obscure einsum bug
# This case revealed a bug in nditer where it reported a stride
# as 'fixed' (0) when it was in fact not fixed during processing
# (0 or 4). The reason for the bug was that the check for a fixed
# stride was using the information from the 2D inner loop reuse
# to restrict the iteration dimensions it had to validate to be
# the same, but that 2D inner loop reuse logic is only triggered
# during the buffer copying step, and hence it was invalid to
# rely on those values. The fix is to check all the dimensions
# of the stride in question, which in the test case reveals that
# the stride is not fixed.
#
# NOTE: This test is triggered by the fact that the default buffersize,
# used by einsum, is 8192, and 3*2731 = 8193, is larger than that
# and results in a mismatch between the buffering and the
# striding for operand A.
A = np.arange(2*3).reshape(2,3).astype(np.float32)
B = np.arange(2*3*2731).reshape(2,3,2731).astype(np.int16)
es = np.einsum('cl,cpx->lpx', A, B)
tp = np.tensordot(A, B, axes=(0, 0))
assert_equal(es, tp)
# The following is the original test case from the bug report,
# made repeatable by changing random arrays to aranges.
A = np.arange(3*3).reshape(3,3).astype(np.float64)
B = np.arange(3*3*64*64).reshape(3,3,64,64).astype(np.float32)
es = np.einsum ('cl,cpxy->lpxy', A,B)
tp = np.tensordot(A,B, axes=(0,0))
assert_equal(es, tp)
def energy(self,R1,M0):
R0=np.tensordot(M0,self.MPO[0],axes=(0,0))
R0=np.tensordot(R0,np.conj(M0),axes=(2,0))
R0=np.transpose(R0,[0,2,4,1,3,5])
energy1=np.squeeze(np.tensordot(R0,R1,axes=([3,4,5],[0,1,2])))
norm=np.tensordot(M0,np.conj(M0),axes=([0,1,2],[0,1,2]))
return energy1/norm
def tensordot_adjoint_0(B, G, axes, A_ndim, B_ndim):
# The adjoint of the operator
# A |--> np.tensordot(A, B, axes)
if B_ndim == 0:
return G * B
G_axes = onp.arange(onp.ndim(G))
if type(axes) is int:
axes = max(axes, 0)
B_axes = onp.arange(B_ndim)
return onp.tensordot(G, B, [G_axes[A_ndim-axes:], B_axes[axes:]])
elif type(axes[0]) is int:
axes = [axes[0] % A_ndim, axes[1] % B_ndim]
B_axes = onp.arange(B_ndim)
return onp.tensordot(G, B, [G_axes[A_ndim-1:], onp.delete(B_axes, axes[1])])
else:
A_axes = onp.arange(A_ndim)
B_axes = onp.arange(B_ndim)
summed_axes = [onp.asarray(axes[0]) % A_ndim,
onp.asarray(axes[1]) % B_ndim]
other_axes = [onp.delete(A_axes, summed_axes[0]),
onp.delete(B_axes, summed_axes[1])]
out = onp.tensordot(G, B, [G_axes[len(other_axes[0]):], other_axes[1]])
perm = onp.argsort(onp.concatenate(
(other_axes[0], summed_axes[0][onp.argsort(summed_axes[1])])))
return onp.transpose(out, perm)
def tensordot_adjoint_1(A, G, axes, A_ndim, B_ndim):
# The adjoint of the operator
# B |--> np.tensordot(A, B, axes)
if A_ndim == 0:
return G * A
G_axes = onp.arange(onp.ndim(G))
if type(axes) is int:
axes = max(axes, 0)
A_axes = onp.arange(A_ndim)
return onp.tensordot(A, G, [A_axes[:A_ndim-axes], G_axes[:A_ndim-axes]])
elif type(axes[0]) is int:
axes = [axes[0] % A_ndim, axes[1] % B_ndim]
A_axes = onp.arange(A_ndim)
return onp.tensordot(A, G, [onp.delete(A_axes, axes[0]), G_axes[:A_ndim-1]])
else:
A_axes = onp.arange(A_ndim)
B_axes = onp.arange(B_ndim)
summed_axes = [onp.asarray(axes[0]) % A_ndim,
onp.asarray(axes[1]) % B_ndim]
other_axes = [onp.delete(A_axes, summed_axes[0]),
onp.delete(B_axes, summed_axes[1])]
out = onp.tensordot(A, G, [other_axes[0], G_axes[:len(other_axes[0])]])
perm = onp.argsort(onp.concatenate(
(summed_axes[1][onp.argsort(summed_axes[0])], other_axes[1])))
return onp.transpose(out, perm)
def H_chol(self):
H = np.tensordot(self.DWt(),self.Wc(),axes=(1,0))
H = np.transpose(H, (0, 2, 1))
H = np.tensordot(self.Wr(), H, axes=(0,0))
H = H.reshape((self.rank_c * self.rank_r, self.rank_c * self.rank_r), order='F')
H+= np.eye(self.rank_r * self.rank_c)
return la.cholesky(H).T
def learn(self, Xd):
# Initialize a data particle
X = [Xd]+[self.X[l]*0+self.O[l]
for l in range(1, len(self.X))]
# Alternate gibbs sampler on data and free particles
for l in (range(1, len(self.X), 2)+range(2, len(self.X), 2))*5:
self.gibbs(X, l)
for l in (range(1, len(self.X), 2)+range(0, len(self.X), 2))*1:
self.gibbs(self.X, l)
# Parameter update
self.W[0] += lr*(np.tensordot(X[0]-self.O[0], X[1]-self.O[1], axes=([0],[0])) -
np.tensordot(self.X[0]-self.O[0], self.X[1]-self.O[1], axes=([0],[0])))/len(Xd)
for i in range(1, len(self.W)):
self.W[i] += lr*(numpy.dot((X[i]-self.O[i]).T, X[i+1]-self.O[i+1]) -
numpy.dot((self.X[i]-self.O[i]).T, self.X[i+1]-self.O[i+1]))/len(Xd)
for i in range(0, len(self.B)):
self.B[i] += lr*(X[i]-self.X[i]).mean(axis=0)
# Reparameterization
for l in range(0, len(self.B)):
self.reparamB(X, l)
for l in range(0, len(self.O)):
self.reparamO(X, l)
def test1():
import axis as ax
import numpy as np
a = ax.Axis((7,), np.ndarray)
ia = ax.IrrepAxis("C2v", (4,0,1,2), np.ndarray)
A = Array(a*a*ia*ia)
B = Array(ia*ia*a*a)
a = np.random.rand(7,7,7,7)
b = np.random.rand(7,7,7,7)
c = np.tensordot(a, b, axes=((2,3),(0,1)))
it_a = zip(range(1),(slice(7),),(slice(7),))
it_ia = zip(range(4),(slice(4),slice(0),slice(1),slice(2)),(slice(0,4),slice(4,4),slice(4,5),slice(5,7)))
for h1, i1, j1 in it_a:
for h2, i2, j2 in it_a:
for h3, i3, j3 in it_ia:
for h4, i4, j4 in it_ia:
A[h1,h2,h3,h4][i1,i2,i3,i4] = a[j1,j2,j3,j4]
B[h3,h4,h1,h2][i3,i4,i1,i2] = b[j3,j4,j1,j2]
col_axes = A.multiaxis.axes[:2]
trc_axes = A.multiaxis.axes[2:]
row_axes = B.multiaxis.axes[2:]
C = Array(col_axes + row_axes)
for col in np.prod(col_axes).iter_array_keytups():
for row in np.prod(row_axes).iter_array_keytups():
for trc in MultiAxis(trc_axes).iter_array_keytups():
if not (0 in A[col+trc].shape or 0 in B[trc+row].shape):
C[col+row] += np.tensordot(A[col+trc], B[trc+row], axes=((2,3),(0,1)))
print np.linalg.norm(C[0,0,0,0] - c)
def concprof(n0, T, r, phi):
kb = 8.61733034e-5
beta = 1.0 / (kb * T)
ux = strain(r, phi)
st = np.zeros((3, 3))
for i in xrange(3):
for j in xrange(3):
for k in xrange(3):
for l in xrange(3):
st[i, j] = Cijkl[i, j, k, l] * ux[k, l] + st[i, j]
F1 = -float(np.tensordot(dpl1, ux))
F2 = -float(np.tensordot(dpl2, ux))
chi1 = n0 / (1 - n0) * np.exp(-F1 * beta)
chi2 = n0 / (1 - n0) * np.exp(-F2 * beta)
n1 = chi1 / (1 + chi1)
n2 = chi2 / (1 + chi2)
n1 = n0 / (np.exp(F1 * beta) * (1 - n0) + n0)
n2 = n0 / (np.exp(F2 * beta) * (1 - n0) + n0)
dn = (n1 + n2) - n0
return dn
def __init__(self, A = None, eps = 1e-14):
if A is None:
self.core = 0
self.u = [0, 0, 0]
self.n = [0, 0, 0]
self.r = [0, 0, 0]
return
N1, N2, N3 = A.shape
B1 = np.reshape(A, (N1, -1), order='F')
B2 = np.reshape(np.transpose(A, [1, 0, 2]), (N2, -1), order='F' )
B3 = np.reshape(np.transpose(A, [2, 0, 1]), (N3, -1), order='F')
U1, V1, r1 = svd_trunc(B1, eps)
U2, V2, r2 = svd_trunc(B2, eps)
U3, V3, r3 = svd_trunc(B3, eps)
G = np.tensordot(A, np.conjugate(U3), (2,0))
G = np.transpose(G, [2, 0, 1])
G = np.tensordot(G, np.conjugate(U2), (2,0))
G = np.transpose(G, [0, 2, 1])
G = np.tensordot(G, np.conjugate(U1), (2,0))
G = np.transpose(G, [2, 1, 0])
self.n = [N1, N2, N3]
self.r = G.shape
self.u = [U1, U2, U3]
self.core = G
def test_solve_fredholm_reconstr_ac():
"""
here we see that the reconstruction quality is independent of the integration weights
differences occur when checking validity of the interpolated time continuous Fredholm equation
"""
_WC_ = 2
def lac(t):
return np.exp(- np.abs(t) - 1j*_WC_*t)
t_max = 10
tol = 2e-10
for ng in range(11,500,30):
t, w = sp.method_kle.get_mid_point_weights_times(t_max, ng)
r = lac(t.reshape(-1,1)-t.reshape(1,-1))
_eig_val, _eig_vec = sp.method_kle.solve_hom_fredholm(r, w)
_eig_vec_ast = np.conj(_eig_vec) # (N_gp, N_ev)
tmp = _eig_val.reshape(1, -1) * _eig_vec # (N_gp, N_ev)
recs_bcf = np.tensordot(tmp, _eig_vec_ast, axes=([1], [1]))
rd = np.max(np.abs(recs_bcf - r) / np.abs(r))
assert rd < tol, "rd={} >= {}".format(rd, tol)
t, w = sp.method_kle.get_simpson_weights_times(t_max, ng)
r = lac(t.reshape(-1, 1) - t.reshape(1, -1))
_eig_val, _eig_vec = sp.method_kle.solve_hom_fredholm(r, w)
_eig_vec_ast = np.conj(_eig_vec) # (N_gp, N_ev)
tmp = _eig_val.reshape(1, -1) * _eig_vec # (N_gp, N_ev)
recs_bcf = np.tensordot(tmp, _eig_vec_ast, axes=([1], [1]))
rd = np.max(np.abs(recs_bcf - r) / np.abs(r))
assert rd < tol, "rd={} >= {}".format(rd, tol)
def eval_basis(self, maps):
"""
It returns basis functions evaluated at quadrature points and mapped
at reference element according to real element.
"""
if self.eval_method == ['grad']:
val = nm.tensordot(self.bfref, maps.inv_jac, axes=(-1, 0))
return val
elif self.eval_method == ['val']:
return self.bfref
elif self.eval_method == ['div']:
val = nm.tensordot(self.bfref, maps.inv_jac, axes=(-1, 0))
val = nm.atleast_3d(nm.einsum('ijkk', val))
return val
elif self.eval_method == ['grad', 'sym', 'Man']:
val = nm.tensordot(self.bfref, maps.inv_jac, axes=(-1, 0))
from sfepy.terms.terms_general import proceed_methods
val = proceed_methods(val, self.eval_method[1:])
return val
else:
msg = "Improper method '%s' for evaluation of basis functions" \
% (self.eval_method)
raise NotImplementedError(msg)
def real(a): # doubled ranks!
b = tensor()
b.n = a.n
b.u[0] = np.concatenate((np.real(a.u[0]), np.imag(a.u[0])), 1)
b.u[1] = np.concatenate((np.real(a.u[1]), np.imag(a.u[1])), 1)
b.u[2] = np.concatenate((np.real(a.u[2]), np.imag(a.u[2])), 1)
R1 = np.zeros((2*a.r[0], a.r[0]), dtype = np.complex128)
R2 = np.zeros((2*a.r[1], a.r[1]), dtype = np.complex128)
R3 = np.zeros((2*a.r[2], a.r[2]), dtype = np.complex128)
R1[:a.r[0], :] = np.identity(a.r[0])
R1[a.r[0]:, :] = 1j*np.identity(a.r[0])
R2[:a.r[1], :] = np.identity(a.r[1])
R2[a.r[1]:, :] = 1j*np.identity(a.r[1])
R3[:a.r[2], :] = np.identity(a.r[2])
R3[a.r[2]:, :] = 1j*np.identity(a.r[2])
GG = np.tensordot(np.transpose(a.core,[2,1,0]),np.transpose(R1), (2,0))
GG = np.tensordot(np.transpose(GG,[0,2,1]),np.transpose(R2), (2,0))
GG = np.transpose(GG,[1,2,0])
b.core = np.real(np.tensordot(GG,np.transpose(R3), (2,0)))
b.r = b.core.shape
return b
def J(theta, hidden_layer_size, num_labels, feat, y, lamb):
"""cost function"""
# Reshape theta into matrices.
# The i stands for 'internal' to indicate that these
# variables are internal to this function.
Theta1i = theta[:hidden_layer_size * feat.shape[1]].reshape( \
hidden_layer_size, feat.shape[1])
Theta2i = theta[hidden_layer_size * feat.shape[1]:].reshape( \
num_labels, hidden_layer_size+1)
# compute the probabilities h by going through the three layers
# layer 1, the input layer
a = feat.T
# layer 2, the hidden layer
z = np.dot(Theta1i,a)
a = g(z)
a = np.vstack((np.ones(a.shape[1]),a))
# layer 3, the output layer
z = np.dot(Theta2i,a)
a = g(z)
h = a.T
# compute the cost
m = feat.shape[0]
cost = -1. / m * (np.tensordot(y, np.log(h)) \
+ np.tensordot(1-y, np.log(1-h)))
# add regularization
cost += lamb * .5 / m * (np.sum(Theta1i[:,1:]**2) \
+ np.sum(Theta2i[:,1:]**2))
return cost
def backward_cpu(self, inputs, grad_outputs):
x, W = inputs[:2]
b = inputs[2] if len(inputs) == 3 else None
if not type_check.same_types(*inputs):
if b is not None:
raise ValueError('numpy and cupy must not be used together\n'
'type(W): {0}, type(x): {1}, type(b): {2}'
.format(type(W), type(x), type(b)))
else:
raise ValueError('numpy and cupy must not be used together\n'
'type(W): {0}, type(x): {1}'
.format(type(W), type(x)))
gy = grad_outputs[0]
kh, kw = W.shape[2:]
col = conv.im2col_cpu(
gy, kh, kw, self.sy, self.sx, self.ph, self.pw)
gW = numpy.tensordot(
x, col, ([0, 2, 3], [0, 4, 5])).astype(W.dtype, copy=False)
gx = numpy.tensordot(
col, W, ([1, 2, 3], [1, 2, 3])).astype(x.dtype, copy=False)
gx = numpy.rollaxis(gx, 3, 1)
if b is None:
return gx, gW
else:
gb = gy.sum(axis=(0, 2, 3))
return gx, gW, gb
def rot_symm(symm_grp_ax, bp_norms_go1, file_path):
"""
Returns the symmetrically equivalent boundary plane normals
Parameters
----------
symm_grp_ax: principle axes for the bicrystal fundamental zone
* numpy array of size (3 x 3)
bp_norms_go1: normalized boundary plane normals in the po1 reference frame
* numpy array of size (n x 3)
file_path: path to the relevant symmetry operations containing pickle file
* string
Returns
-------
symm_bpn_go1: symmetrically equivalent boundary plane normals in the po1 reference frame
* numpy array of size (m x n x 3); m == order of bicrystal point group symmetry group
Notes
------
"""
bpn_rot = np.dot(np.linalg.inv(symm_grp_ax), bp_norms_go1.transpose()).transpose()
symm_mat = pickle.load(open(file_path, 'rb'))
### np.dot returns the sum product of the last axis of the first matrix with the second to last axis of the second
### advisable to use np.tensordot instead to avoid confusion !!
symm_bpn_rot_gop1 = np.tensordot(symm_mat, bpn_rot.transpose(), 1).transpose((1, 2, 0))
symm_bpn_go1 = np.tensordot(symm_grp_ax, symm_bpn_rot_gop1, 1).transpose(2, 1, 0)
return symm_bpn_go1
def princ2(cauchyStress,C):
"""
Calculate the principal values using
Sig = C:T:s form
Arguments
---------
cauchyStress - the linearly transformed stress Sigma
C (6x6) tensor for linear transformation
Returns
-------
S1
S2
S3
"""
sqrt=np.sqrt
s = cauchyStress.copy()
T = cpb_lib.returnT()
Sig = np.tensordot(C,np.tensordot(T,s,axes=(1,0)),axes=(1,0))
Sxx,Syy,Szz,Sxy = Sig[0], Sig[1], Sig[2], Sig[5]
S1 = 0.5*(Sxx + Syy + sqrt( (Sxx-Syy)**2 + 4*Sxy**2))
S2 = 0.5*(Sxx + Syy - sqrt( (Sxx-Syy)**2 + 4*Sxy**2))
S3 = Szz
return S1, S2, S3
def get_chisquareds(self, n):
dd = self.get_datum(n)
if n == 0:
foo = np.sum(np.sum(np.tensordot(dd, self.get_ivarts(), axes=(1,1)) * dd[:,None,:],
axis=2), axis=0)
return np.sum(np.sum(np.tensordot(dd, self.get_ivarts(), axes=(1,1)) * dd[:,None,:],
axis=2), axis=0) # should be length T
def update(self):
check_inputs_synchronized(self)
y_dot = self.input.y_dot
y = self.input.y
u = self.input.u
T = self.T
check_multiple([
('shape(x),(array[HxW]|array[N])', y),
('shape(x),(array[HxW]|array[N])', y_dot),
('array[K]', u),
('array[Kx2xHxW]|array[Kx1xN]', T), # TODO: use references
])
K = u.size
# update covariance of gradients
gy = generalized_gradient(y)
Tgy = np.tensordot([1, 1], T * gy , axes=(0, 1))
y_dot_pred = np.tensordot(u, Tgy, axes=(0, 0))
assert y_dot_pred.ndim == 2
error = np.maximum(0, -y_dot_pred * y_dot)
# error = np.abs(np.sign(y_dot_pred) - np.sign(y_dot))
self.output.y_dot_pred = y_dot_pred
self.output.error = error
self.output.error_sensel = np.sum(error, axis=0)
def four_index_transform_cholesky(ao_integrals, orb0, orb1=None, method='tensordot'):
"""Perform four index transformation on a Cholesky-decomposed four-index object.
Parameters
----------
oa_integrals : np.ndarray, shape=(nvec, nbasis, nbasis)
Cholesky decomposition of four-index object in the AO basis.
orb0
A Orbitals object with molecular orbitals.
orb1
Can be provided to transform the second index differently.
method
Either ``einsum`` or ``tensordot`` (default).
"""
if orb1 is None:
orb1 = orb0
result = np.zeros(ao_integrals.shape)
if method == 'einsum':
result = np.einsum('ai,kac->kic', orb0.coeffs, ao_integrals)
result = np.einsum('cj,kic->kij', orb1.coeffs, result)
elif method == 'tensordot':
result = np.tensordot(ao_integrals, orb0.coeffs, axes=([1],[0]))
result = np.tensordot(result, orb1.coeffs, axes=([1],[0]))
else:
raise ValueError('The method must either be \'einsum\' or \'tensordot\'.')
return result
def dev_trapz(ds): d = ds.variables['development'][:]/1e4 t = ds.variables['time'][:] t_floor = np.floor(t).astype(int) #now set up a loop to solve the index for each whole integer or iterate up #we can figure out if we have a leap year or not also by looking at the max day #i.e. np.max(t_floor) == 365 , means this is a leap year dataset max_days = np.max(t_floor) + 1 z,m,n = d.shape daily_dev = np.zeros((max_days, m,n)) start = 0 #first integration point k_array = np.zeros(max_days) for k in range(max_days): ones = np.ones((m,n)) if k == max_days-1: tm = np.tensordot(t[start:], ones, axes=0) daily_dev[k] = np.trapz(d[start:],tm,axis=0) else: end = np.argmax(t_floor>k) tm = np.tensordot(t[start:end+1], ones, axes=0) #plus one to integrate the end point daily_dev[k] = np.trapz(d[start:end+1],tm,axis=0) #define the new start start = end k_array[k] = k+1 #counter for each day return(daily_dev, k_array)
def discretizedGaussian(amp, mu, cov, grid):
'''Convenience method for discretized Gaussian evaluation'''
#eigenvalue decomposition of precision matrix
P = la.inv(cov) #precision matrix
evl, M = la.eig(P)
#assert np.allclose(np.diag(evl), iM.dot(cov).dot(M))
#check if covariance is positive definite
if np.any(evl < 0):
raise ValueError('Covariance matrix should be positive definite')
#make column vector for arithmetic
mu = np.array(mu, ndmin=grid.ndim, dtype=float).T
evl = np.array(evl, ndmin=grid.ndim, dtype=float).T
xm = grid - mu #(2,...) shape
#return M, xm
f = np.sqrt(2 / evl)
pf = np.sqrt(np.pi / evl)
td0 = np.tensordot(M, xm + 0.5, 1)
td1 = np.tensordot(M, xm - 0.5, 1)
w = pf*(erf(f*td0) - erf(f*td1))
return amp * np.prod(w, axis=0)
def d2y_SS(z_i):
DF,df,Hf = self.DF(z_i),self.df(z_i),self.Hf(z_i)
d = self.get_d(z_i)
DFi = DF[n:]
return np.tensordot(np.linalg.inv(DFi[:,y] + DFi[:,e].dot(df) + DFi[:,v].dot(Ivy)),
-self.HFhat[S,S](z_i) - np.tensordot(DFi[:,e],quadratic_dot(Hf,d[y,S],d[y,S]),1)
, axes=1)
def expect(mps,opl): #opl is a list of matrices including identities.but this is not necessary #contract s first expect=np.array([[1.]]) for i in range(len(mps.Ms)): overlap=np.tensordot(mps.Ms[i].conjugate().transpose(),np.tensordot(expect,mps.Ms[i],1),1) expect=np.tensordot(opl[i],overlap,axes=([0,1],[1,2])) return float(expect)
def right_sweep_projection(flmps0,flmps1,qtarget,thresh,Dcut,debug,\
ifBlockSingletEmbedding=False):
nsite, ne, ms, sval = qtarget
print('\n[mpo_dmrg_samps.right_sweep_projection] (nsite,,ne,ms,sval)=',
(nsite, ne, ms, sval))
flmps1['nsite'] = nsite
# Reduced qnums - the basis states are ordered to maximize the efficiency.
qnumsN = numpy.array([[0., 0., 1.], [1., 0.5, 1.], [2., 0., 1.]])
qnums1 = [None] * (nsite + 1)
# This works perfectly.
qnums1[nsite] = numpy.array([[0., 0., 1.]])
wmat = numpy.identity(1)
for isite in range(nsite - 1, -1, -1):
site0 = mpo_dmrg_io.loadSite(flmps0, isite, False)
# *** Change the direction for combinations
site0 = numpy.einsum('lnr->rnl', site0)
# Expand |M> basis to |SM> basis
wA0 = numpy.tensordot(wmat, site0, axes=([1], [0])) # Ll,lNr->LNr
# Transform to combined basis
qnumsL, wA1 = util_tensor.spinCouple(wA0, qnums1[isite + 1], qnumsN)
# Quasi-RDM (reduced): each dim is {[(SaSb)S,na*nb]}
classes, qrdmL = util_tensor.quasiRDM(wA1, qnumsL)
if debug:
print('trace=', numpy.trace(qrdmL), qrdmL.shape, len(classes))
# Decimation: [(SaSb)S,na*nb]=>[S,r] for all possible Sa,Sb
# Therefore, rotL is a matrix with dimensions [(SaSb)S,na*nb]*[S,r]!
dwts, qred, rotL, sigs = mpo_dmrg_qparser.rdm_blkdiag(
qrdmL, classes, thresh, Dcut, debug)
#
# | (0) w[i-1] = <l[i-1]|m[i-1]>: expansion coeff of |l(NM)> to |m(NSM)>
# |-----|------| (1) update formula for w[i]
# | L | (2) L[i] - isometry
# | | | (3) L[i]*t[CG] - new site A[i]
# | t[CG] |
# | / \ | |
# | w---A[i]---A[i+1]---
# |------------|
#
# Screen
qnumsl, rotL = util_spinsym.symScreen(isite,
nsite,
rotL,
qred,
sigs,
ne,
sval,
debug,
status='R')
# A[lnr]
site1 = util_tensor.expandRotL(rotL, qnums1[isite + 1], qnumsN, qnumsL,
qnumsl)
# srotR = rotL^\dagger * wA
srotR = numpy.tensordot(site1, wA0, axes=([0, 1], [0,
1])) # LNR,LNr->Rr
# Print:
print(' ---> isite=',isite,'site0=',site0.shape,'rotL=',rotL.shape,\
'site1=',site1.shape) #,'srotR=',srotR.shape
print()
# Check left canonical form
if debug:
tmp = numpy.tensordot(site1, site1,
axes=([0, 1], [0, 1])) # lna,lnb->ab
d = tmp.shape[0]
diff = numpy.linalg.norm(tmp - numpy.identity(d))
print(' diff=', diff)
assert diff < 1.e-10
# Save
if isite > 0:
wmat = srotR.copy()
qnums1[isite] = qnumsl.copy()
flmps1.create_dataset('rotL' + str(isite),
data=rotL,
compression='lzf')
# *** Change the direction back [Right canonical form]
site1 = numpy.einsum('rnl->lnr', site1)
flmps1.create_dataset('site' + str(isite),
data=site1,
compression='lzf')
else:
# In case of singlet embedding, the in bond of the first site
# is expanded in the subroutine expandRotL in the left sweep for nonsinglet state.
if ifBlockSingletEmbedding and abs(sval) > 1.e-10:
im = int(ms + sval)
srotR = srotR[:, im].reshape(srotR.shape[0], 1)
# Special treatment for the last site by invoking symmetry selection!
info, qnumsl, rotL, site1, wt = util_tensor.lastSite(
rotL, site1, srotR, qnumsl, ne, ms, sval)
if info:
qnums1[0] = qnumsl.copy()
flmps1.create_dataset('rotL' + str(isite),
data=rotL,
compression='lzf')
site1 = numpy.einsum('rnl->lnr', site1)
flmps1.create_dataset('site' + str(isite),
data=site1,
compression='lzf')
# Dump on flmps1
if info:
print('\nfinalize right_sweep ...')
qnumsm = [None] * (nsite + 1)
for isite in range(1, nsite + 1):
qnumsm[isite] = util_spinsym.expandSM(qnums1[isite])
# Right case
qnumsm[0] = numpy.array([[ne, ms]])
# DUMP qnums
for isite in range(nsite + 1):
flmps1['qnum' + str(isite)] = qnumsm[isite]
flmps1['qnumNS' + str(isite)] = qnums1[isite]
# Check
for isite in range(nsite + 1):
dimr = util_spinsym.dim_red(qnums1[isite])
print(' ibond/dim(NS)/dim(NSM)=', isite, dimr, len(qnumsm[isite]))
return wt
def tensordot(self, a: Tensor, b: Tensor, axes: Sequence[Sequence[int]]):
return np.tensordot(a, b, axes)
def outer_product(self, tensor1: Tensor, tensor2: Tensor) -> Tensor:
return np.tensordot(tensor1, tensor2, 0)
def min_tens(tens, rmax=10, nswp=10, verb=True, smooth_fun=None):
"""Find (approximate) minimal element in a TT-tensor."""
if smooth_fun is None:
smooth_fun = lambda p, lam: (math.pi / 2 - np.arctan(p - lam))
d = tens.d
Rx = [[]] * (d + 1) # Python list for the interfaces
Rx[0] = np.ones((1, 1))
Rx[d] = np.ones((1, 1))
Jy = [np.empty(0, dtype=np.int)] * (d + 1)
ry = rmax * np.ones(d + 1, dtype=np.int)
ry[0] = 1
ry[d] = 1
n = tens.n
elements_seen = 0
phi_left = [np.empty(0)] * (d + 1)
phi_left[0] = np.array([1])
phi_right = [np.empty(0)] * (d + 1)
phi_right[d] = np.array([1])
cores = tt.tensor.to_list(tens)
# Fill initial multiindex J randomly.
grid = [np.reshape(range(n[i]), (n[i], 1)) for i in xrange(d)]
for i in xrange(d - 1):
ry[i + 1] = min(ry[i + 1], n[i] * ry[i])
ind = sorted(np.random.permutation(ry[i] * n[i])[0:ry[i + 1]])
w1 = mkron(np.ones((n[i], 1), dtype=np.int), Jy[i])
w2 = mkron(grid[i], np.ones((ry[i], 1), dtype=np.int))
Jy[i + 1] = np.hstack((w1, w2))
Jy[i + 1] = reshape(Jy[i + 1], (ry[i] * n[i], -1))
Jy[i + 1] = Jy[i + 1][ind, :]
phi_left[i + 1] = np.tensordot(phi_left[i], cores[i], 1)
phi_left[i + 1] = reshape(phi_left[i + 1], (ry[i] * n[i], -1))
phi_left[i + 1] = phi_left[i + 1][ind, :]
swp = 0
dirn = -1
i = d - 1
lm = float('Inf')
while swp < nswp:
# Right-to-left sweep
# The idea: compute the current core; compute the function of it;
# Shift locally or globally? Local shift would be the first try
# Compute the current core
if np.size(Jy[i]) == 0:
w1 = np.zeros((ry[i] * n[i] * ry[i + 1], 0), dtype=np.int)
else:
w1 = mkron(np.ones((n[i] * ry[i + 1], 1), dtype=np.int), Jy[i])
w2 = mkron(mkron(np.ones((ry[i + 1], 1), dtype=np.int), grid[i]),
np.ones((ry[i], 1), dtype=np.int))
if np.size(Jy[i + 1]) == 0:
w3 = np.zeros((ry[i] * n[i] * ry[i + 1], 0), dtype=np.int)
else:
w3 = mkron(Jy[i + 1], np.ones((ry[i] * n[i], 1), dtype=np.int))
J = np.hstack((w1, w2, w3))
phi_right[i] = np.tensordot(cores[i], phi_right[i + 1], 1)
phi_right[i] = reshape(phi_right[i], (-1, n[i] * ry[i + 1]))
cry = np.tensordot(phi_left[i],
np.tensordot(cores[i], phi_right[i + 1], 1), 1)
elements_seen += cry.size
cry = reshape(cry, (ry[i], n[i], ry[i + 1]))
min_cur = np.min(cry.flatten("F"))
ind_cur = np.argmin(cry.flatten("F"))
if lm > min_cur:
lm = min_cur
x_full = J[ind_cur, :]
val = tens[x_full]
if verb:
print 'New record:', val, 'Point:', x_full, 'elements seen:', elements_seen
cry = smooth_fun(cry, lm)
if dirn < 0 and i > 0:
cry = reshape(cry, (ry[i], n[i] * ry[i + 1]))
cry = cry.T
#q, r = np.linalg.qr(cry)
u, s, v = mysvd(cry, full_matrices=False)
ry[i] = min(ry[i], rmax)
q = u[:, :ry[i]]
ind = rect_maxvol(q)[0] # maxvol(q)
ry[i] = ind.size
w1 = mkron(np.ones((ry[i + 1], 1), dtype=np.int), grid[i])
if np.size(Jy[i + 1]) == 0:
w2 = np.zeros((n[i] * ry[i + 1], 0), dtype=np.int)
else:
w2 = mkron(Jy[i + 1], np.ones((n[i], 1), dtype=np.int))
Jy[i] = np.hstack((w1, w2))
Jy[i] = reshape(Jy[i], (n[i] * ry[i + 1], -1))
Jy[i] = Jy[i][ind, :]
phi_right[i] = np.tensordot(cores[i], phi_right[i + 1], 1)
phi_right[i] = reshape(phi_right[i], (-1, n[i] * ry[i + 1]))
phi_right[i] = phi_right[i][:, ind]
if dirn > 0 and i < d - 1:
cry = reshape(cry, (ry[i] * n[i], ry[i + 1]))
q, r = np.linalg.qr(cry)
#ind = maxvol(q)
ind = rect_maxvol(q)[0]
ry[i + 1] = ind.size
phi_left[i + 1] = np.tensordot(phi_left[i], cores[i], 1)
phi_left[i + 1] = reshape(phi_left[i + 1], (ry[i] * n[i], -1))
phi_left[i + 1] = phi_left[i + 1][ind, :]
w1 = mkron(np.ones((n[i], 1), dtype=np.int), Jy[i])
w2 = mkron(grid[i], np.ones((ry[i], 1), dtype=np.int))
Jy[i + 1] = np.hstack((w1, w2))
Jy[i + 1] = reshape(Jy[i + 1], (ry[i] * n[i], -1))
Jy[i + 1] = Jy[i + 1][ind, :]
i += dirn
if i == d or i == -1:
dirn = -dirn
i += dirn
swp = swp + 1
return val, x_full
def rhs_default_PS(z, t=0):
x, y = z
q = np.array([(x - x_1), (y - y_1)])
return np.tensordot(A, q, axes=[(1, ), (0, )])
def neuron_output(weights, inputs):
if weights.ndim == 2:
return sigmoid(weights.dot(inputs))
elif weights.ndim == 3:
return sigmoid(np.tensordot(weights, inputs))
def left_sweep_projection(flmps0,flmps1,qtarget,thresh,Dcut,debug,\
ifBlockSingletEmbedding=False,\
ifBlockSymScreen=False,\
ifpermute=False):
nsite, ne, ms, sval = qtarget
print('\n[mpo_dmrg_samps.left_sweep_projection] (nsite,ne,ms,sval)=',(nsite,ne,ms,sval),\
'ifBlockSE=',ifBlockSingletEmbedding)
flmps1['nsite'] = nsite
# Reduced qnums - the basis states are ordered to maximize the efficiency.
qnumsN = numpy.array([[0., 0., 1.], [1., 0.5, 1.], [2., 0., 1.]])
qnums1 = [None] * (nsite + 1)
# Do not use singlet embedding for the first purification sweep,
# otherwise, it will mess up with other spins.
if ifBlockSingletEmbedding == False:
qnums1[0] = numpy.array([[0., 0., 1.]])
wmat = numpy.identity(1)
# Quantum numbers
ne_eff = ne
sval_eff = sval
ms_eff = ms
else:
qnums1[0] = numpy.array([[2 * sval, sval, 1.]]) # (N,S)=(2*S,S)
ndeg = int(2 * sval + 1)
wmat = numpy.zeros((ndeg, 1))
# Coupled to a noninteracting state |S(-M)>
# to create a broken symetry state |S(-M)>*|SM>
# which has the singlet component |00>/sqrt(2*S+1).
im = int(-ms + sval)
wmat[im] = 1.0
# Quantum numbers
ne_eff = ne + 2.0 * sval
sval_eff = 0.0
ms_eff = 0.0
# Start conversion for the following sites
for isite in range(nsite):
site0 = mpo_dmrg_io.loadSite(flmps0, isite, False)
# Expand |M> basis to |SM> basis
wA0 = numpy.tensordot(wmat, site0, axes=([1], [0])) # Ll,lNr->LNr
# Transform to combined basis
if ifpermute and ifBlockSingletEmbedding and isite == 0:
wA0 = wA0.transpose(1, 0, 2)
sgn = numpy.array(
[1., (-1.0)**(int(2 * sval)), (-1.0)**(int(2 * sval)), 1.])
wA0 = numpy.einsum('n,nvr->nvr', sgn, wA0)
qnumsL, wA1 = util_tensor.spinCouple(wA0, qnumsN, qnums1[isite])
else:
qnumsL, wA1 = util_tensor.spinCouple(wA0, qnums1[isite], qnumsN)
# Quasi-RDM (reduced): each dim is {[(SaSb)S,na*nb]}
classes, qrdmL = util_tensor.quasiRDM(wA1, qnumsL)
if debug:
print('trace=', numpy.trace(qrdmL), qrdmL.shape, len(classes))
# Decimation: [(SaSb)S,na*nb]=>[S,r] for all possible Sa,Sb
# Therefore, rotL is a matrix with dimensions [(SaSb)S,na*nb]*[S,r]!
if ifBlockSingletEmbedding and isite == 0:
dwts, qred, rotL, sigs = mpo_dmrg_qparser.rdm_blkdiag(
qrdmL, classes, -1.e-10, -1, debug)
else:
dwts, qred, rotL, sigs = mpo_dmrg_qparser.rdm_blkdiag(
qrdmL, classes, thresh, Dcut, debug)
#
# | (0) w[i-1] = <l[i-1]|m[i-1]>: expansion coeff of |l(NM)> to |m(NSM)>
# |-----|------| (1) update formula for w[i]
# | L | (2) L[i] - isometry
# | | | (3) L[i]*t[CG] - new site A[i]
# | t[CG] |
# | / \ | |
# | w---A[i]---A[i+1]---
# |------------|
#
# Screen
qnumsl,rotL = util_spinsym.symScreen(isite,nsite,rotL,qred,sigs,ne_eff,sval_eff,debug,\
ifBlockSymScreen=ifBlockSymScreen)
if ifpermute and ifBlockSingletEmbedding and isite == 0:
# A[lnr]
site1 = util_tensor.expandRotL(rotL, qnumsN, qnums1[isite], qnumsL,
qnumsl)
# srotR = rotL^\dagger * wA
srotR = numpy.tensordot(site1, wA0,
axes=([0, 1], [0, 1])) # LNR,LNr->Rr
site1 = site1.transpose(1, 0, 2)
else:
# A[lnr]
site1 = util_tensor.expandRotL(rotL, qnums1[isite], qnumsN, qnumsL,
qnumsl)
# srotR = rotL^\dagger * wA
srotR = numpy.tensordot(site1, wA0,
axes=([0, 1], [0, 1])) # LNR,LNr->Rr
# Print:
print(' ---> isite=',isite,'site0=',site0.shape,'rotL=',rotL.shape,\
'site1=',site1.shape) #,'srotR=',srotR.shape
print()
# Check left canonical form
if debug:
tmp = numpy.tensordot(site1, site1,
axes=([0, 1], [0, 1])) # lna,lnb->ab
d = tmp.shape[0]
diff = numpy.linalg.norm(tmp - numpy.identity(d))
print(' diff=', diff)
assert diff < 1.e-10
# Save
if isite < nsite - 1:
wmat = srotR.copy()
qnums1[isite + 1] = qnumsl.copy()
flmps1.create_dataset('rotL' + str(isite),
data=rotL,
compression='lzf')
flmps1.create_dataset('site' + str(isite),
data=site1,
compression='lzf')
else:
# Special treatment for the last site by invoking symmetry selection!
info, qnumsl, rotL, site1, wt = util_tensor.lastSite(
rotL, site1, srotR, qnumsl, ne_eff, ms_eff, sval_eff)
if info:
qnums1[nsite] = qnumsl.copy()
flmps1.create_dataset('rotL' + str(isite),
data=rotL,
compression='lzf')
flmps1.create_dataset('site' + str(isite),
data=site1,
compression='lzf')
# Dump on flmps1
if info:
print('\nfinalize left_sweep ...')
qnumsm = [None] * (nsite + 1)
for isite in range(nsite):
qnumsm[isite] = util_spinsym.expandSM(qnums1[isite])
# Left case
qnumsm[nsite] = numpy.array([[ne, ms]])
# DUMP qnums
for isite in range(nsite + 1):
flmps1['qnum' + str(isite)] = qnumsm[isite]
flmps1['qnumNS' + str(isite)] = qnums1[isite]
# Check
for isite in range(nsite + 1):
dimr = util_spinsym.dim_red(qnums1[isite])
print(' ibond/dim(NS)/dim(NSM)=', isite, dimr, len(qnumsm[isite]))
return wt
def form_compressed_hamiltonian_offdiag_1block_diff(vecs_l, vecs_r, Hi, Hij,
differences):
# {{{
"""
Find (1-site) Hamiltonian matrix in between differently compressed spaces:
i.e.,
<Abcd| H(1) |ab'c'd'> = <A|Ha|a> * del(bb') * del(cc')
or like,
<aBcd| H(1) |abcd> = <B|Hb|b> * del(aa') * del(cc')
Notice that the full Hamiltonian will also have the two-body part:
<Abcd| H(2) |abcd> = <Ab|Hab|ab'> del(cc') + <Ac|Hac|ac'> del(bb')
differences = [blocks which are not diagonal]
i.e., for <Abcd|H(1)|abcd>, differences = [1]
for <Abcd|H(1)|aBcd>, differences = [1,2], and all values are zero for H(1)
vecs_l [vecs_A, vecs_B, vecs_C, ...]
vecs_r [vecs_A, vecs_B, vecs_C, ...]
"""
dim_l = 1 # dimension of left subspace
dim_r = 1 # dimension of right subspace
dim_same = 1 # dimension of space spanned by all those fragments is the same space (i.e., not the blocks inside differences)
dims_l = [
] # list of mode dimensions (size of hilbert space on each fragment)
dims_r = [
] # list of mode dimensions (size of hilbert space on each fragment)
block_curr = differences[0] # current block which is offdiagonal
# vecs_l correspond to the bra states
# vecs_r correspond to the ket states
assert (len(vecs_l) == len(vecs_r))
for bi, b in enumerate(vecs_l):
dim_l = dim_l * b.shape[1]
dims_l.extend([b.shape[1]])
if bi != block_curr:
dim_same *= b.shape[1]
dim_same_check = 1
for bi, b in enumerate(vecs_r):
dim_r = dim_r * b.shape[1]
dims_r.extend([b.shape[1]])
if bi != block_curr:
dim_same_check *= b.shape[1]
assert (dim_same == dim_same_check)
H = np.zeros((dim_l, dim_r))
print " Size of Hamitonian block: ", H.shape
assert (len(dims_l) == len(dims_r))
n_dims = len(dims_l)
H1 = cp.deepcopy(Hi)
H2 = cp.deepcopy(Hij)
## Rotate the single block Hamiltonians into the requested single site basis
#for vi in range(0,n_dims):
# l = vecs_l[vi]
# r = vecs_r[vi]
# H1[vi] = l.T.dot(Hi[vi]).dot(r)
# Rotate the double block Hamiltonians into the appropriate single site basis
# for vi in range(0,n_dims):
# for wi in range(0,n_dims):
# if wi>vi:
# vw_l = np.kron(vecs_l[vi],vecs_l[wi])
# vw_r = np.kron(vecs_r[vi],vecs_r[wi])
# H2[(vi,wi)] = vw_l.T.dot(Hij[(vi,wi)]).dot(vw_r)
dimsdims = np.append(
dims_l, dims_r
) # this is the tensor layout for the many-body Hamiltonian in the current subspace
vecs = vecs_l
dims = dims_l
dim = dim_l
# Add up all the one-body contributions, making sure that the results is properly dimensioned for the
# target subspace
dim_i1 = 1 # dimension of space for fragments to the left of the current 'different' fragment
dim_i2 = 1 # dimension of space for fragments to the right of the current 'different' fragment
# <abCdef|H1|abcdef> = eye(a) x eye(b) x <C|H1|c> x eye(d) x eye(e) x eye(f)
for vi in range(n_dims):
if vi < block_curr:
dim_i1 *= dims_l[vi]
assert (dims_l[vi] == dims_r[vi])
elif vi > block_curr:
dim_i2 *= dims_l[vi]
assert (dims_l[vi] == dims_r[vi])
# Rotate the current single block Hamiltonian into the requested single site basis
l = vecs_l[block_curr]
r = vecs_r[block_curr]
h1_block = l.T.dot(Hi[block_curr]).dot(r)
i1 = np.eye(dim_i1)
i2 = np.eye(dim_i2)
H += np.kron(i1, np.kron(h1_block, i2))
# <abCdef|H2(0,2)|abcdef> = <aC|H2|ac> x eye(b) x eye(d) x eye(e) x eye(f) = <aCbdef|H2|acbdef>
#
# then transpose:
# flip Cb and cb
H.shape = (dims_l + dims_r)
for bi in range(0, block_curr):
#print "block_curr, bi", block_curr, bi
vw_l = np.kron(vecs_l[bi], vecs_l[block_curr])
vw_r = np.kron(vecs_r[bi], vecs_r[block_curr])
h2 = vw_l.T.dot(Hij[(bi, block_curr)]).dot(
vw_r
) # i.e. get reference to <aC|Hij[0,2]|a'c>, where block_curr = 2 and bi = 0
dim_i = dim_same / dims[bi]
#i1 = np.eye(dim_i)
#tmp_h2 = np.kron( h2, i1)
h2.shape = (dims_l[bi], dims_l[block_curr], dims_r[bi],
dims_r[block_curr])
tens_inds = []
tens_inds.extend([bi, block_curr])
tens_inds.extend([bi + n_dims, block_curr + n_dims])
for bbi in range(0, n_dims):
if bbi != bi and bbi != block_curr:
tens_inds.extend([bbi])
tens_inds.extend([bbi + n_dims])
#print "h2.shape", h2.shape,
h2 = np.tensordot(h2, np.eye(dims[bbi]), axes=0)
#print "h2.shape", h2.shape,
sort_ind = np.argsort(tens_inds)
H += h2.transpose(sort_ind)
#print "tens_inds", tens_inds
#print "sort_ind", sort_ind
#print "h2", h2.transpose(sort_ind).shape
#H += h2
for bi in range(block_curr, n_dims):
if bi == block_curr:
continue
#print "block_curr, bi", block_curr, bi
vw_l = np.kron(vecs_l[block_curr], vecs_l[bi])
vw_r = np.kron(vecs_r[block_curr], vecs_r[bi])
h2 = vw_l.T.dot(Hij[(block_curr, bi)]).dot(
vw_r
) # i.e. get reference to <aC|Hij[0,2]|a'c>, where block_curr = 2 and bi = 0
h2.shape = (dims_l[block_curr], dims_l[bi], dims_r[block_curr],
dims_r[bi])
tens_inds = []
tens_inds.extend([block_curr, bi])
tens_inds.extend([block_curr + n_dims, bi + n_dims])
for bbi in range(0, n_dims):
if bbi != bi and bbi != block_curr:
tens_inds.extend([bbi])
tens_inds.extend([bbi + n_dims])
#print "h2.shape", h2.shape,
h2 = np.tensordot(h2, np.eye(dims[bbi]), axes=0)
#print "h2.shape", h2.shape,
sort_ind = np.argsort(tens_inds)
#print "h2", h2.transpose(sort_ind).shape
#print "tens_inds", tens_inds
#print "sort_ind", sort_ind
H += h2.transpose(sort_ind)
#print "h2.shape", h2.shape
H.shape = (dim_l, dim_r)
return H
def main(options):
CF.CreatePipeOutput(options.DestDir + '/' + options.Logfile)
#VC.OptionsCheck(options,'Phonons')
CF.PrintMainHeader('Bandstructures', vinfo, options)
fdf = glob.glob(options.FCwildcard +
'/RUN.fdf') # This should be made an input flag
SCDM = Supercell_DynamicalMatrix(fdf, accuracy=options.accuracy)
# Write high-symmetry path
WritePath(options.DestDir + '/symmetry-path', SCDM.Sym.path, options.steps)
# Write mesh
k1, k2, k3 = eval(options.mesh)
rvec = 2 * N.pi * N.array([SCDM.Sym.b1, SCDM.Sym.b2, SCDM.Sym.b3])
import Kmesh
# Full mesh
kmesh = Kmesh.kmesh(2**k1,
2**k2,
2**k3,
meshtype=['LIN', 'LIN', 'LIN'],
invsymmetry=False)
WriteKpoints(options.DestDir + '/mesh_%ix%ix%i' % tuple(kmesh.Nk),
N.dot(kmesh.k, rvec))
# Mesh reduced by inversion symmetry
kmesh = Kmesh.kmesh(2**k1,
2**k2,
2**k3,
meshtype=['LIN', 'LIN', 'LIN'],
invsymmetry=True)
WriteKpoints(options.DestDir + '/mesh_%ix%ix%i_invsym' % tuple(kmesh.Nk),
N.dot(kmesh.k, rvec))
# Evaluate electron k-points
if options.kfile:
# Prepare Hamiltonian etc in Gamma for whole supercell
natoms = SIO.GetFDFlineWithDefault(fdf[0], 'NumberOfAtoms', int, -1,
'Error')
SCDM.PrepareGradients(options.onlySdir,
N.array([0., 0., 0.]),
1,
natoms,
AbsEref=False,
atype=N.complex)
SCDM.nao = SCDM.h0.shape[-1]
SCDM.FirstOrb = SCDM.OrbIndx[0][0] # First atom = 1
SCDM.LastOrb = SCDM.OrbIndx[SCDM.Sym.basis.NN -
1][1] # Last atom = Sym.NN
SCDM.rednao = SCDM.LastOrb + 1 - SCDM.FirstOrb
# Read kpoints
kpts, dk, klabels, kticks = ReadKpoints(options.kfile)
if klabels:
# Only write ascii if labels exist
WriteKpoints(options.DestDir + '/kpoints', kpts, klabels)
# Prepare netcdf
ncfn = options.DestDir + '/Electrons.nc'
ncf = NC4.Dataset(ncfn, 'w')
# Grid
ncf.createDimension('gridpts', len(kpts))
ncf.createDimension('vector', 3)
grid = ncf.createVariable('grid', 'd', ('gridpts', 'vector'))
grid[:] = kpts
grid.units = '1/Angstrom'
# Geometry
ncf.createDimension('atoms', SCDM.Sym.basis.NN)
xyz = ncf.createVariable('xyz', 'd', ('atoms', 'vector'))
xyz[:] = SCDM.Sym.basis.xyz
xyz.units = 'Angstrom'
pbc = ncf.createVariable('pbc', 'd', ('vector', 'vector'))
pbc.units = 'Angstrom'
pbc[:] = [SCDM.Sym.a1, SCDM.Sym.a2, SCDM.Sym.a3]
rvec1 = ncf.createVariable('rvec', 'd', ('vector', 'vector'))
rvec1.units = '1/Angstrom (incl. factor 2pi)'
rvec1[:] = rvec
ncf.sync()
# Loop over kpoints
for i, k in enumerate(kpts):
if i < 100: # Print only for the first 100 points
ev, evec = SCDM.ComputeElectronStates(k, verbose=True)
else:
ev, evec = SCDM.ComputeElectronStates(k, verbose=False)
# otherwise something simple
if i % 100 == 0:
print '%i out of %i k-points computed' % (i, len(kpts))
if i == 0:
ncf.createDimension('nspin', SCDM.nspin)
ncf.createDimension('orbs', SCDM.rednao)
ncf.createDimension('bands', SCDM.rednao)
evals = ncf.createVariable('eigenvalues', 'd',
('gridpts', 'nspin', 'bands'))
evals.units = 'eV'
evecsRe = ncf.createVariable(
'eigenvectors.re', 'd',
('gridpts', 'nspin', 'orbs', 'bands'))
evecsIm = ncf.createVariable(
'eigenvectors.im', 'd',
('gridpts', 'nspin', 'orbs', 'bands'))
# Check eigenvectors
print 'SupercellPhonons: Checking eigenvectors at', k
for j in range(SCDM.nspin):
ev2 = N.diagonal(
MM.mm(MM.dagger(evec[j]), SCDM.h0_k[j], evec[j]))
print ' ... spin %i: Allclose=' % j, N.allclose(ev[j],
ev2,
atol=1e-5,
rtol=1e-3)
ncf.sync()
# Write to NetCDF
evals[i, :] = ev
evecsRe[i, :] = evec.real
evecsIm[i, :] = evec.imag
ncf.sync()
# Include basis orbitals in netcdf file
lasto = SCDM.OrbIndx[:SCDM.Sym.basis.NN + 1, 0]
orbbasis = SIO.BuildBasis(fdf[0], 1, SCDM.Sym.basis.NN, lasto)
# Note that the above basis is for the geometry with an atom FC-moved in z.
#print dir(orbbasis)
#print orbbasis.xyz # Hence, this is not the correct geometry of the basis atoms!
center = ncf.createVariable('orbcenter', 'i', ('orbs', ))
center[:] = N.array(orbbasis.ii - 1, dtype='int32')
center.description = 'Atom index (counting from 0) of the orbital center'
nn = ncf.createVariable('N', 'i', ('orbs', ))
nn[:] = N.array(orbbasis.N, dtype='int32')
ll = ncf.createVariable('L', 'i', ('orbs', ))
ll[:] = N.array(orbbasis.L, dtype='int32')
mm = ncf.createVariable('M', 'i', ('orbs', ))
mm[:] = N.array(orbbasis.M, dtype='int32')
# Cutoff radius and delta
Rc = ncf.createVariable('Rc', 'd', ('orbs', ))
Rc[:] = orbbasis.coff
Rc.units = 'Angstrom'
delta = ncf.createVariable('delta', 'd', ('orbs', ))
delta[:] = orbbasis.delta
delta.units = 'Angstrom'
# Radial components of the orbitals
ntb = len(orbbasis.orb[0])
ncf.createDimension('ntb', ntb)
rii = ncf.createVariable('rii', 'd', ('orbs', 'ntb'))
rii[:] = N.outer(orbbasis.delta, N.arange(ntb))
rii.units = 'Angstrom'
radialfct = ncf.createVariable('radialfct', 'd', ('orbs', 'ntb'))
radialfct[:] = orbbasis.orb
# Sort eigenvalues to connect crossing bands?
if options.sorting:
for i in range(SCDM.nspin):
evals[:, i, :] = SortBands(evals[:, i, :])
# Produce nice plots if labels exist
if klabels:
if SCDM.nspin == 1:
PlotElectronBands(options.DestDir + '/Electrons.agr', dk,
evals[:, 0, :], kticks)
elif SCDM.nspin == 2:
PlotElectronBands(options.DestDir + '/Electrons.UP.agr', dk,
evals[:, 0, :], kticks)
PlotElectronBands(options.DestDir + '/Electrons.DOWN.agr', dk,
evals[:, 1, :], kticks)
ncf.close()
# Compute phonon eigenvalues
SCDM.SymmetrizeFC(options.radius)
SCDM.SetMasses()
if options.qfile:
qpts, dq, qlabels, qticks = ReadKpoints(options.qfile)
else:
qpts, dq, qlabels, qticks = ReadKpoints(options.DestDir +
'/symmetry-path')
if qlabels:
# Only write ascii if labels exist
WriteKpoints(options.DestDir + '/qpoints', qpts, qlabels)
# Prepare netcdf
ncfn = options.DestDir + '/Phonons.nc'
ncf = NC4.Dataset(ncfn, 'w')
# Grid
ncf.createDimension('gridpts', len(qpts))
ncf.createDimension('vector', 3)
grid = ncf.createVariable('grid', 'd', ('gridpts', 'vector'))
grid[:] = qpts
grid.units = '1/Angstrom'
# Geometry
ncf.createDimension('atoms', SCDM.Sym.basis.NN)
xyz = ncf.createVariable('xyz', 'd', ('atoms', 'vector'))
xyz[:] = SCDM.Sym.basis.xyz
xyz.units = 'Angstrom'
pbc = ncf.createVariable('pbc', 'd', ('vector', 'vector'))
pbc.units = 'Angstrom'
pbc[:] = [SCDM.Sym.a1, SCDM.Sym.a2, SCDM.Sym.a3]
rvec1 = ncf.createVariable('rvec', 'd', ('vector', 'vector'))
rvec1.units = '1/Angstrom (incl. factor 2pi)'
rvec1[:] = rvec
ncf.sync()
# Loop over q
for i, q in enumerate(qpts):
if i < 100: # Print only for the first 100 points
hw, U = SCDM.ComputePhononModes_q(q, verbose=True)
else:
hw, U = SCDM.ComputePhononModes_q(q, verbose=False)
# otherwise something simple
if i % 100 == 0:
print '%i out of %i q-points computed' % (i, len(qpts))
if i == 0:
ncf.createDimension('bands', len(hw))
ncf.createDimension('displ', len(hw))
evals = ncf.createVariable('eigenvalues', 'd',
('gridpts', 'bands'))
evals.units = 'eV'
evecsRe = ncf.createVariable('eigenvectors.re', 'd',
('gridpts', 'bands', 'displ'))
evecsIm = ncf.createVariable('eigenvectors.im', 'd',
('gridpts', 'bands', 'displ'))
# Check eigenvectors
print 'SupercellPhonons.Checking eigenvectors at', q
tmp = MM.mm(N.conjugate(U), SCDM.FCtilde, N.transpose(U))
const = PC.hbar2SI * (1e20 / (PC.eV2Joule * PC.amu2kg))**0.5
hw2 = const * N.diagonal(tmp)**0.5 # Units in eV
print ' ... Allclose=', N.allclose(hw,
N.absolute(hw2),
atol=1e-5,
rtol=1e-3)
ncf.sync()
evals[i] = hw
evecsRe[i] = U.real
evecsIm[i] = U.imag
ncf.sync()
# Sort eigenvalues to connect crossing bands?
if options.sorting:
evals = SortBands(evals)
# Produce nice plots if labels exist
if qlabels:
PlotPhononBands(options.DestDir + '/Phonons.agr', dq,
N.array(evals[:]), qticks)
ncf.close()
# Compute e-ph couplings
if options.kfile and options.qfile:
SCDM.ReadGradients(AbsEref=False)
ncf = NC4.Dataset(options.DestDir + '/EPH.nc', 'w')
ncf.createDimension('kpts', len(kpts))
ncf.createDimension('qpts', len(qpts))
ncf.createDimension('modes', len(hw))
ncf.createDimension('nspin', SCDM.nspin)
ncf.createDimension('bands', SCDM.rednao)
ncf.createDimension('vector', 3)
kgrid = ncf.createVariable('kpts', 'd', ('kpts', 'vector'))
kgrid[:] = kpts
qgrid = ncf.createVariable('qpts', 'd', ('qpts', 'vector'))
qgrid[:] = qpts
evalfkq = ncf.createVariable('evalfkq', 'd',
('kpts', 'qpts', 'nspin', 'bands'))
# First (second) band index n (n') is the initial (final) state, i.e.,
# Mkq(k,q,mode,spin,n,n') := < n',k+q | dV_q(mode) | n,k >
MkqAbs = ncf.createVariable(
'Mkqabs', 'd',
('kpts', 'qpts', 'modes', 'nspin', 'bands', 'bands'))
GkqAbs = ncf.createVariable(
'Gkqabs', 'd',
('kpts', 'qpts', 'modes', 'nspin', 'bands', 'bands'))
ncf.sync()
# Loop over k-points
for i, k in enumerate(kpts):
kpts[i] = k
# Compute initial electronic states
evi, eveci = SCDM.ComputeElectronStates(k, verbose=True)
# Loop over q-points
for j, q in enumerate(qpts):
# Compute phonon modes
hw, U = SCDM.ComputePhononModes_q(q, verbose=True)
# Compute final electronic states
evf, evecf = SCDM.ComputeElectronStates(k + q, verbose=True)
evalfkq[i, j, :] = evf
# Compute electron-phonon couplings
m, g = SCDM.ComputeEPHcouplings_kq(
k, q) # (modes,nspin,bands,bands)
# Data to file
# M (modes,spin,i,l) = m(modes,k,j) init(i,j) final(k,l)
# 0 1 2 0,1 0 1
# ^-------^
# ^----------------------^
for ispin in range(SCDM.nspin):
evecfd = MM.dagger(evecf[ispin]) # (bands,bands)
M = N.tensordot(N.tensordot(m[:, ispin],
eveci[ispin],
axes=[2, 0]),
evecfd,
axes=[1, 1])
G = N.tensordot(N.tensordot(g[:, ispin],
eveci[ispin],
axes=[2, 0]),
evecfd,
axes=[1, 1])
MkqAbs[i, j, :, ispin] = N.absolute(M)
GkqAbs[i, j, :, ispin] = N.absolute(G)
ncf.sync()
ncf.close()
return SCDM.Sym.path
def form_compressed_hamiltonian_offdiag_2block_diff(vecs_l, vecs_r, Hi, Hij,
differences):
# {{{
"""
Find Hamiltonian matrix in between differently compressed spaces:
i.e.,
<Abcd| H(0,2) |a'b'C'd'> = <Ac|Ha|a'C'> * del(bb')
differences = [blocks which are not diagonal]
i.e., for <Abcd|H(1)|abCd>, differences = [0,2]
vecs_l [vecs_A, vecs_B, vecs_C, ...]
vecs_r [vecs_A, vecs_B, vecs_C, ...]
"""
dim_l = 1 # dimension of left subspace
dim_r = 1 # dimension of right subspace
dim_same = 1 # dimension of space spanned by all those fragments is the same space (i.e., not the blocks inside differences)
dims_l = [
] # list of mode dimensions (size of hilbert space on each fragment)
dims_r = [
] # list of mode dimensions (size of hilbert space on each fragment)
assert (
len(differences) == 2
) # make sure we are not trying to get H(1) between states with multiple fragments orthogonal
block_curr1 = differences[0] # current block which is offdiagonal
block_curr2 = differences[1] # current block which is offdiagonal
# vecs_l correspond to the bra states
# vecs_r correspond to the ket states
assert (len(vecs_l) == len(vecs_r))
for bi, b in enumerate(vecs_l):
dim_l = dim_l * b.shape[1]
dims_l.extend([b.shape[1]])
if bi != block_curr1 and bi != block_curr2:
dim_same *= b.shape[1]
dim_same_check = 1
for bi, b in enumerate(vecs_r):
dim_r = dim_r * b.shape[1]
dims_r.extend([b.shape[1]])
if bi != block_curr1 and bi != block_curr2:
dim_same_check *= b.shape[1]
assert (dim_same == dim_same_check)
H = np.zeros((dim_l, dim_r))
print " Size of Hamitonian block: ", H.shape
assert (len(dims_l) == len(dims_r))
n_dims = len(dims_l)
dimsdims = np.append(
dims_l, dims_r
) # this is the tensor layout for the many-body Hamiltonian in the current subspace
# Add up all the one-body contributions, making sure that the results is properly dimensioned for the
# target subspace
dim_i1 = 1 # dimension of space for fragments to the left of the current 'different' fragment
dim_i2 = 1 # dimension of space for fragments to the right of the current 'different' fragment
# <Abcdef|H2(0,2)|abCdef> = <Ac|H2|aC> x eye(b) x eye(d) x eye(e) x eye(f) = <aCbdef|H2|acbdef>
#
# then transpose:
# flip Cb and cb
H.shape = (dims_l + dims_r)
assert (block_curr1 < block_curr2)
#print " block_curr1, block_curr2", block_curr1, block_curr2
vw_l = np.kron(vecs_l[block_curr1], vecs_l[block_curr2])
vw_r = np.kron(vecs_r[block_curr1], vecs_r[block_curr2])
h2 = vw_l.T.dot(Hij[(block_curr1, block_curr2)]).dot(
vw_r
) # i.e. get reference to <aC|Hij[0,2]|a'c>, where block_curr = 2 and bi = 0
h2.shape = (dims_l[block_curr1], dims_l[block_curr2], dims_r[block_curr1],
dims_r[block_curr2])
tens_inds = []
tens_inds.extend([block_curr1, block_curr2])
tens_inds.extend([block_curr1 + n_dims, block_curr2 + n_dims])
for bbi in range(0, n_dims):
if bbi != block_curr1 and bbi != block_curr2:
tens_inds.extend([bbi])
tens_inds.extend([bbi + n_dims])
#print "h2.shape", h2.shape,
assert (dims_l[bbi] == dims_r[bbi])
dims = dims_l
h2 = np.tensordot(h2, np.eye(dims[bbi]), axes=0)
#print "h2.shape", h2.shape,
sort_ind = np.argsort(tens_inds)
H += h2.transpose(sort_ind)
#print "tens_inds", tens_inds
#print "sort_ind", sort_ind
#print "h2", h2.transpose(sort_ind).shape
#H += h2
H.shape = (dim_l, dim_r)
return H
def contract(self, tensor2, indexPair):
data = np.tensordot(self.data, tensor2.data, axes=indexPair)
return numpyBackend(data=data)
def gre_sim(T1, T2, TR=12e-3, TE=6e-3, alpha=np.pi/3, field_map=None, phi=0,
dphi=0, M0=1, tol=1e-5, maxiter=None, spoil=True):
'''Simulate GRE pulse sequence.
Parameters
==========
T1 : array_like
longitudinal exponential decay time constant.
T2 : array_like
Transverse exponential decay time constant.
TR : float, optional
repetition time.
TE : float, optional
echo time.
alpha : float, optional
flip angle.
field_map : array_like, optional
offresonance field map (in hertz).
phi : float, optional
Reference starting phase.
dphi : float, optional
phase cycling of RF pulses.
M0 : array_like, optional
proton density.
tol : float, optional
Maximum difference between voxel intensity iter to iter until stop.
maxiter : int, optional
number of excitations till steady state.
Returns
=======
array_like
Complex transverse magnetization (Mx + 1j*My)
Notes
=====
maxiter=None will run until difference between all voxel intensities
iteration to iteration is within given tolerance, tol (default=1e-5).
'''
if not isinstance(T1, np.ndarray):
T1 = np.array([T1])
T2 = np.array([T2])
M0 = np.array([M0])
if field_map is None:
field_map = np.zeros(T1.shape)
# We have 2 states: current, previous iteration
Mgre = np.zeros((3, 2) + T1.shape)
Mgre[2, 0, ...] = M0
# Rotation matrix each Rx pulse
rxalpha = np.array([
[1, 0, 0],
[0, np.cos(alpha), np.sin(alpha)],
[0, -np.sin(alpha), np.cos(alpha)]])
# first flip
c_phi, s_phi = np.cos(phi), np.sin(phi)
rzdphi = np.array([
[c_phi, s_phi, 0],
[-s_phi, c_phi, 0],
[0, 0, 1]])
rznegdphi = np.array([
[c_phi, -s_phi, 0],
[s_phi, c_phi, 0],
[0, 0, 1]])
rot_vec = np.linalg.multi_dot(
(rzdphi, rxalpha, rznegdphi))
Mgre[:, 0, ...] = np.tensordot(rot_vec, Mgre[:, 0, ...], axes=1)
# Precompute some values we each loop, making sure we don't divide by zero
idx1 = np.nonzero(T1)
idx2 = np.nonzero(T2)
E1 = np.zeros(T1.shape)
E2 = np.zeros(T2.shape)
E1[idx1] = np.exp(-TR/T1[idx1])
E2[idx2] = np.exp(-TR/T2[idx2])
cycles = field_map*TR
rotation_angle = np.fmod(cycles, 1)*2*np.pi
c_ra = np.cos(rotation_angle)
s_ra = np.sin(rotation_angle)
def iter_fun(Mgre, phi=0):
'''Heavy lifting function run each iteration'''
# relaxation
Mgre[0, 1, ...] = Mgre[0, 0, ...]*E2 # x
Mgre[1, 1, ...] = Mgre[1, 0, ...]*E2 # y
Mgre[2, 1, ...] = 1 + (Mgre[2, 0, ...] - 1)*E1 # z
# Here's where we spend most of our time:
for idx, _fm in np.ndenumerate(field_map):
rzoffres = np.array([
[c_ra[idx], s_ra[idx], 0],
[-s_ra[idx], c_ra[idx], 0],
[0, 0, 1]])
Mgre[(slice(None), 1) + idx] = rzoffres.dot(
Mgre[(slice(None), 1) + idx])
# next tip - delete phase information! to make it gre
if spoil:
Mgre[0, 1, ...] = 0
Mgre[1, 1, ...] = 0
# will got over 2*pi but shouldn't matter
phi += dphi
c_phi, s_phi = np.cos(phi), np.sin(phi)
rzdphi = np.array([
[c_phi, s_phi, 0],
[-s_phi, c_phi, 0],
[0, 0, 1]])
rznegdphi = np.array([
[c_phi, -s_phi, 0],
[s_phi, c_phi, 0],
[0, 0, 1]])
rot_vec = np.linalg.multi_dot((rzdphi, rxalpha, rznegdphi))
Mgre[:, 1, ...] = np.tensordot(rot_vec, Mgre[:, 1, ...], axes=1)
# Update prev iteration
Mgre[:, 0, ...] = Mgre[:, 1, ...]
return(Mgre, phi)
# Do either fixed number of iter or until tolerance achieved
if maxiter is not None:
# assume steady state after iter flips
for _n in trange(maxiter, leave=False, desc='GRE steady-state'):
Mgre, phi = iter_fun(Mgre, phi)
else:
# Run until all voxels within tolerance
Mgre_prev = np.ones(Mgre.shape)*np.inf
stop_cond = False
while not stop_cond:
Mgre, phi = iter_fun(Mgre, phi)
stop_cond = np.any(np.abs(Mgre - Mgre_prev) < tol)
Mgre_prev = Mgre
# take steady state sample and relax in TE
Mss = np.zeros((3,) + T1.shape)
E2 = np.zeros(T2.shape)
E2[idx2] = np.exp(-TE/T2[idx2])
Mss[0, ...] = Mgre[0, -1, ...]*E2
Mss[1, ...] = Mgre[1, -1, ...]*E2
Mss[2, ...] = Mgre[2, -1, ...]
cycles = field_map*TE
rotation_angle = np.fmod(cycles, 1)*2*np.pi
c_ra, s_ra = np.cos(rotation_angle), np.sin(rotation_angle)
for idx, _fm in np.ndenumerate(field_map):
rzoffres = np.array([
[c_ra[idx], s_ra[idx], 0],
[-s_ra[idx], c_ra[idx], 0],
[0, 0, 1]])
Mss[(slice(None),) + idx] = rzoffres.dot(Mss[(slice(None),) + idx])
return(Mss[0, ...] + 1j*Mss[1, ...])
def get_corr_pred(self, sctx, eps_app_eng):
'''
Corrector predictor computation.
@param eps_app_eng input variable - engineering strain
'''
# -----------------------------------------------------------------------------------------------
# check if average strain is to be used for damage evaluation
# -----------------------------------------------------------------------------------------------
# if eps_avg != None:
# pass
# else:
# eps_avg = eps_app_eng
# -----------------------------------------------------------------------------------------------
# for debugging purposes only: if elastic_debug is switched on, linear elastic material is used
# -----------------------------------------------------------------------------------------------
if self.elastic_debug:
# NOTE: This must be copied otherwise self.D2_e gets modified when
# essential boundary conditions are inserted
D2_e = copy(self.D2_e)
sig_eng = tensordot(D2_e, eps_app_eng, [[1], [0]])
return sig_eng, D2_e
# print 'sctx_correc', sctx.shape
# -----------------------------------------------------------------------------------------------
# update state variables
# -----------------------------------------------------------------------------------------------
# if sctx.update_state_on:
# #eps_n = eps_avg - d_eps
# e_max = self._get_state_variables(sctx, eps_app_eng)
# sctx.mats_state_array[:] = e_max
#----------------------------------------------------------------------
# if the regularization using the crack-band concept is on calculate the
# effective element length in the direction of principle strains
#----------------------------------------------------------------------
# if self.regularization:
# h = self.get_regularizing_length(sctx, eps_app_eng)
# self.phi_fn.h = h
#------------------------------------------------------------------
# Damage tensor (2th order):
#------------------------------------------------------------------
phi_ij = self._get_phi_mtx(sctx, eps_app_eng)
# print 'phi_ij', phi_ij
#------------------------------------------------------------------
# Damage tensor (4th order) using product- or sum-type symmetrization:
#------------------------------------------------------------------
beta_ijkl = self._get_beta_tns(phi_ij)
# print 'beta_ijkl', beta_ijkl
# print 'beta_ijkl', beta_ijkl
#------------------------------------------------------------------
# Damaged stiffness tensor calculated based on the damage tensor beta4:
#------------------------------------------------------------------
# (cf. [Jir99] Eq.(7): C = beta * D_e * beta^T),
# minor symmetry is tacitly assumed ! (i.e. beta_ijkl = beta_jilk)
# D4_mdm = tensordot(
# beta_ijkl, tensordot(self.D4_e, beta_ijkl, [[2, 3], [2, 3]]),
# [[2, 3], [0, 1]])
D4_mdm_ijmn = einsum('ijkl,klsr,mnsr->ijmn', beta_ijkl, self.D4_e,
beta_ijkl)
# print self.D4_e
# print 'D4_mdm_ijmn', D4_mdm_ijmn
# print 'Damged stiffness', D4_mdm_ijmn
# print self.D4_e
#------------------------------------------------------------------
# Reduction of the fourth order tensor to a matrix assuming minor and major symmetry:
#------------------------------------------------------------------
#D2_mdm = self.map_tns4_to_tns2(D4_mdm_ijmn)
# print'D2_mdm', D2_mdm.shape
#----------------------------------------------------------------------
# Return stresses (corrector) and damaged secant stiffness matrix (predictor)
#----------------------------------------------------------------------
eps_e_mtx = eps_app_eng
# print'eps_e_mtx', eps_e_mtx.shape
#sig_eng = tensordot(D2_mdm, eps_e_mtx, [[1], [0]])
sig_eng = einsum('ijmn,mn -> ij', D4_mdm_ijmn, eps_e_mtx)
# print 'sig_eng', sig_eng
return sig_eng, D4_mdm_ijmn
def perturb_continuation_params(x,
equil_obj,
deltas,
args,
pert_order=1,
verbose=False,
timer=None):
"""perturbs an equilibrium wrt the continuation parameters
Parameters
----------
x : ndarray
state vector
equil_obj : function
equilibrium objective function
deltas : ndarray
changes in the continuation parameters
args : tuple
additional arguments passed to equil_obj
pert_order : int
order of perturbation (1=linear, 2=quadratic) (Default value = 1)
verbose : int or bool
level of output to display (Default value = False)
timer : Timer
Timer object (Default value = None)
Returns
-------
x : ndarray
perturbed state vector
timer : Timer
Timer object with timing data
"""
delta_strings = ['boundary', 'pressure', 'zeta'
] if len(deltas) == 3 else [None] * len(deltas)
f = equil_obj(x, *args)
dimF = len(f)
dimX = len(x)
dimC = 1
if timer is None:
timer = Timer()
timer.start('Total perturbation')
# 1st order
if pert_order >= 1:
# partial derivatives wrt x
timer.start('df/dx computation')
obj_jac_x = jac(equil_obj, argnum=0).compute
Jx = obj_jac_x(x, *args).reshape((dimF, dimX))
timer.stop('df/dx computation')
RHS = f
if verbose > 1:
timer.disp('df/dx computation')
# partial derivatives wrt c
for i in range(deltas.size):
if deltas[i] != 0:
if verbose > 1:
print("Perturbing {}".format(delta_strings[i]))
timer.start('df/dc computation ({})'.format(delta_strings[i]))
obj_jac_c = jac(equil_obj, argnum=6 + i).compute
Jc = obj_jac_c(x, *args).reshape((dimF, dimC))
timer.stop("df/dc computation ({})".format(delta_strings[i]))
RHS += np.tensordot(Jc, np.atleast_1d(deltas[i]), axes=1)
if verbose > 1:
timer.disp("df/dc computation ({})".format(
delta_strings[i]))
# 2nd order
if pert_order >= 2:
# partial derivatives wrt x
Jxi = np.linalg.pinv(Jx, rcond=1e-6)
timer.start("df/dxx computation")
obj_jac_xx = jac(jac(equil_obj, argnum=0).compute, argnum=0).compute
Jxx = obj_jac_xx(x, *args).reshape((dimF, dimX, dimX))
timer.stop("df/dxx computation")
RHS += 0.5 * np.tensordot(Jxx,
np.tensordot(np.tensordot(Jxi, RHS, axes=1),
np.tensordot(
RHS.T, Jxi.T, axes=1),
axes=0),
axes=2)
if verbose > 1:
timer.disp("df/dxx computation")
# partial derivatives wrt c
for i in range(deltas.size):
if deltas[i] != 0:
if verbose > 1:
print("Perturbing {}".format(delta_strings[i]))
timer.start("df/dcc computation ({})".format(delta_strings[i]))
obj_jac_cc = jac(jac(equil_obj, argnum=6 + i).compute,
argnum=6 + i).compute
Jcc = obj_jac_cc(x, *args).reshape((dimF, dimC, dimC))
timer.stop("df/dcc computation ({})".format(delta_strings[i]))
RHS += 0.5 * np.tensordot(Jcc,
np.tensordot(
np.atleast_1d(deltas[i]),
np.atleast_1d(deltas[i]),
axes=0),
axes=2)
if verbose > 1:
timer.disp("df/dcc computation ({})".format(
delta_strings[i]))
timer.start("df/dxc computation ({})".format(delta_strings[i]))
obj_jac_xc = jac(jac(equil_obj, argnum=0).compute,
argnum=6 + i).compute
Jxc = obj_jac_xc(x, *args).reshape((dimF, dimX, dimC))
timer.stop("df/dxc computation ({})".format(delta_strings[i]))
RHS -= np.tensordot(Jxc,
np.tensordot(Jxi,
np.tensordot(RHS,
np.atleast_1d(
deltas[i]),
axes=0),
axes=1),
axes=2)
if verbose > 1:
timer.disp("df/dxc computation ({})".format(
delta_strings[i]))
# perturbation
if pert_order > 0:
dx = -np.linalg.lstsq(Jx, RHS, rcond=1e-6)[0]
else:
dx = np.zeros_like(x)
timer.stop('Total perturbation')
if verbose > 1:
timer.disp('Total perturbation')
return x + dx, timer
def Two_Site_Rho(self, l, k):
L = self.sim['L']
chi = self.sim['chi']
d = self.sim['d']
Lambda = self.Lambda
Gamma = self.Gamma
swap = 0
# Assume l < k
if (k < l):
temp = k
k = l
l = temp
swap = 1
# Form tensor:
if (k < L - 2):
tensor_l = np.tensordot(np.diag(Lambda[k]), Gamma[k], axes=(0, 1))
tensor_l = np.tensordot(tensor_l,
np.diag(Lambda[k + 1]),
axes=(-1, 0))
tensor_r = np.tensordot(np.diag(Lambda[k + 1]),
np.conjugate(Gamma[k]),
axes=(0, -1))
tensor_r = np.tensordot(tensor_r, np.diag(Lambda[k]), axes=(-1, 0))
tensor = np.tensordot(tensor_l, tensor_r, axes=(-1, 0))
# Arrange indices as (ik, ik', ak-1, ak-1')
tensor = np.transpose(tensor, (1, 2, 0, 3))
elif (k == L - 2):
tensor_l = np.tensordot(np.diag(Lambda[k]), Gamma[k], axes=(0, 1))
tensor_l = np.tensordot(tensor_l, np.identity(chi), axes=(-1, 0))
tensor_r = np.tensordot(np.identity(chi),
np.conjugate(Gamma[k]),
axes=(0, -1))
tensor_r = np.tensordot(tensor_r, np.diag(Lambda[k]), axes=(-1, 0))
tensor = np.tensordot(tensor_l, tensor_r, axes=(-1, 0))
# Arrange indices as (ik, ik', ak-1, ak-1')
tensor = np.transpose(tensor, (1, 2, 0, 3))
elif (k == L - 1):
tensor = np.outer(Gamma[k], np.conjugate(Gamma[k]))
tensor = np.reshape(tensor, (d, chi, d, chi))
tensor = np.transpose(tensor, (0, 2, 1, 3))
for i in range(0, l - k - 1):
tensor = np.tensordot(Gamma[k - 1], tensor, axes=(-1, 2))
tensor = np.tensordot(np.diag(Lambda[k - 1]), tensor, axes=(-1, 1))
tensor = np.tensordot(tensor,
np.conjugate(Gamma[k - 1]),
axes=([1, -1], [0, -1]))
tensor = np.tensordot(tensor, np.diag(Lambda[k - 1]), axes=(-1, 0))
np.transpose(tensor, (1, 2, 0, 3))
if (l != 0):
rho_ts = np.tensordot(tensor,
np.conjugate(Gamma[l]),
axes=(-1, -1))
rho_ts = np.tensordot(rho_ts, np.diag(Lambda[l]), axes=(-1, 0))
rho_ts = np.tensordot(Gamma[l], rho_ts, axes=(-1, 2))
rho_ts = np.tensordot(np.diag(Lambda[l]),
rho_ts,
axes=([0, 1], [1, -1]))
# Transpose to (il, il', ik, ik')
rho_lk = np.transpose(rho_ts, (0, 3, 1, 2))
else:
rho_ts = np.tensordot(tensor,
np.conjugate(Gamma[l]),
axes=(-1, -1))
rho_ts = np.tensordot(Gamma[l], rho_ts, axes=(-1, 2))
# Transpose to (il, il', ik, ik')
rho_lk = np.transpose(rho_ts, (0, 3, 1, 2))
return rho_lk
def __call__(self, x, idx=slice(None)): """evaluate trajectory at specified time""" theta = (x - self.xold) / self.hout theta1 = 1.0 - theta coef = vstack([ones_like(theta),theta,theta1,theta,theta1]).cumprod(axis=0) return tensordot(coef,self.rcont,(0,0))[newaxis,idx]
def Update(self, V, l):
d = self.sim['d']
L = self.sim['L']
chi = self.sim['chi']
delta = self.sim['delta']
# Build the appropriate Theta tensor
Theta = self.Build_Theta(l)
# Apply the unitary matrix V
Theta = np.tensordot(V, Theta, axes=([2, 3], [0, 1]))
# Need to treat boundary subsystems differently...
if (l != 0 and l != L - 2):
# Reshape to a square matrix and do singular value decomposition:
Theta = np.reshape(np.transpose(Theta, (0, 2, 1, 3)),
(d * chi, d * chi))
Theta *= 1 / np.linalg.norm(np.absolute(Theta))
# A and transpose.C contain the new Gamma[l] and Gamma[l+1]
# B contains new Lambda[l]
A, B, C = np.linalg.svd(Theta)
C = np.transpose(C)
# Truncate at chi eigenvalues and enforce normalization
norm = np.linalg.norm(B[0:chi])
B = B / norm
self.Lambda[l, :] = B[0:chi]
# Keep track of the truncation error accumulated on this step
self.tau += delta * (1 - norm**2)
# Find the new Gammas:
# Gamma_l:
A = np.reshape(A[:, 0:chi], (d, chi, chi))
Gamma_l_new = np.tensordot(A,
np.diag(OneOver(self.Lambda[l - 1])),
axes=(1, 0))
self.Gamma[l] = np.transpose(Gamma_l_new, (0, 2, 1))
# Gamma_(l+1):
C = np.reshape(C[:, 0:chi], (d, chi, chi))
Gamma_lp1_new = np.tensordot(C,
np.diag(OneOver(self.Lambda[l + 1])),
axes=(-1, 0))
self.Gamma[l + 1] = Gamma_lp1_new
# The Gamma_0 tensor has one less index... need to treat slightly differently
elif (l == 0):
# Reshape to a square matrix and do singular value decomposition:
Theta = np.reshape(Theta, (d, d * chi))
Theta *= 1 / np.linalg.norm(np.absolute(Theta))
# A and transpose.C contain the new Gamma[l] and Gamma[l+1]
# B contains new Lambda[l]
A, B, C = np.linalg.svd(Theta)
C = np.transpose(C)
# Enforce normalization
# Don't need to truncate here because chi is bounded by
# the dimension of the smaller subsystem, here equals d < chi
norm = np.linalg.norm(B)
B = B / norm
self.Lambda[l, 0:d] = B
# Keep track of the truncation error accumulated on this step
self.tau += delta * (1 - norm**2)
# Find the new Gammas:
# Gamma_l:
# Note: can't truncate because the matrix
# is smaller than in the non-edge cases
# Right multiplying by Lambda_l means that we're effectively storing
# the A tensors instead of the Gamma tensors.
# See Quantum Gases: Finite Temperature and Non-Equilibrium Dynamics, pg. 341
self.Gamma[l][:, 0:d] = A
# Treat the l = 1 case normally...
C = np.reshape(C[:, 0:chi], (d, chi, chi))
Gamma_lp1_new = np.tensordot(C,
np.diag(OneOver(self.Lambda[l +
1, :])),
axes=(-1, 0))
self.Gamma[l + 1] = Gamma_lp1_new
elif (l == L - 2):
# Reshape to a square matrix and do singular value decomposition:
Theta = np.reshape(np.transpose(Theta, (0, 2, 1)), (d * chi, d))
Theta *= 1 / np.linalg.norm(np.absolute(Theta))
# A and transpose.C contain the new Gamma[l] and Gamma[l+1]
# B contains new Lambda[l]
A, B, C = np.linalg.svd(Theta)
C = np.transpose(C)
# Enforce normalization
# Don't need to truncate here because chi is bounded by
# the dimension of the smaller subsystem, here equals d < chi
norm = np.linalg.norm(B)
B = B / norm
self.Lambda[l, 0:d] = B
# Keep track of the truncation error accumulated on this step
self.tau += delta * (1 - norm**2)
# Find the new Gammas:
# Treat the L-2 case normally:
A = np.reshape(A[:, 0:chi], (d, chi, chi))
# Right multiplying by Lambda_l means that we're effectively storing
# the A tensors instead of the Gamma tensors.
# See Quantum Gases: Finite Temperature and Non-Equilibrium Dynamics, pg. 341
Gamma_l_new = np.tensordot(A,
np.diag(OneOver(self.Lambda[l - 1])),
axes=(1, 0))
self.Gamma[l] = np.transpose(Gamma_l_new, (0, 2, 1))
# Gamma_(L-1):
# Note: can't truncate because the matrix
# is smaller than in the non-edge cases
self.Gamma[l + 1][:, 0:d] = C
def test_np_tensordot():
class TestTensordot(HybridBlock):
def __init__(self, axes):
super(TestTensordot, self).__init__()
self._axes = axes
def hybrid_forward(self, F, a, b):
return F.np.tensordot(a, b, self._axes)
def tensordot_backward(a, b, axes=2):
if (a.ndim < 1) or (b.ndim < 1):
raise ValueError('An input is zero-dim')
if _np.isscalar(axes):
a_axes_summed = [i + a.ndim - axes for i in range(axes)]
b_axes_summed = [i for i in range(axes)]
else:
if len(axes) != 2:
raise ValueError('Axes must consist of two arrays.')
a_axes_summed, b_axes_summed = axes
if _np.isscalar(a_axes_summed):
a_axes_summed = a_axes_summed,
if _np.isscalar(b_axes_summed):
b_axes_summed = b_axes_summed,
for i in range(len(a_axes_summed)):
a_axes_summed[i] = (a_axes_summed[i] + a.ndim) % a.ndim
for i in range(len(b_axes_summed)):
b_axes_summed[i] = (b_axes_summed[i] + b.ndim) % b.ndim
if len(a_axes_summed) != len(b_axes_summed):
raise ValueError('Axes length mismatch')
a_axes_remained = []
for i in range(a.ndim):
if not (i in a_axes_summed):
a_axes_remained.append(i)
a_axes = a_axes_remained[:] + a_axes_summed[:]
b_axes_remained = []
for i in range(b.ndim):
if not (i in b_axes_summed):
b_axes_remained.append(i)
b_axes = b_axes_summed[:] + b_axes_remained[:]
ad1 = _np.prod([a.shape[i] for i in a_axes_remained
]) if len(a_axes_remained) > 0 else 1
ad2 = _np.prod([a.shape[i] for i in a_axes_summed
]) if len(a_axes_summed) > 0 else 1
bd1 = _np.prod([b.shape[i] for i in b_axes_summed
]) if len(b_axes_summed) > 0 else 1
bd2 = _np.prod([b.shape[i] for i in b_axes_remained
]) if len(b_axes_remained) > 0 else 1
out_grad = _np.ones((ad1, bd2))
new_a = _np.transpose(a, a_axes)
new_a_shape = new_a.shape[:]
new_a = new_a.reshape((ad1, ad2))
new_b = _np.transpose(b, b_axes)
new_b_shape = new_b.shape[:]
new_b = new_b.reshape((bd1, bd2))
reverse_a_axes = [0 for i in a_axes]
for i in range(len(a_axes)):
reverse_a_axes[a_axes[i]] = i
reverse_b_axes = [0 for i in b_axes]
for i in range(len(b_axes)):
reverse_b_axes[b_axes[i]] = i
grad_b = _np.dot(new_a.T, out_grad).reshape(new_b_shape)
grad_b = _np.transpose(grad_b, reverse_b_axes)
grad_a = _np.dot(out_grad, new_b.T).reshape(new_a_shape)
grad_a = _np.transpose(grad_a, reverse_a_axes)
return [grad_a, grad_b]
# test non zero size input
tensor_shapes = [
((3, 5), (5, 4), 1), # (a_shape, b_shape, axes)
((3, ), (3, ), 1),
((3, 4, 5, 3, 2), (5, 3, 2, 1, 2), 3),
((3, 5, 4, 3, 2), (2, 3, 5, 1, 2), [[1, 3, 4], [2, 1, 0]]),
((3, 5, 4), (5, 4, 3), [[1, 0, 2], [0, 2, 1]]),
((3, 5, 4), (5, 3, 4), [[2, 0], [-1, -2]]),
((2, 2), (2, 2), 2),
((3, 5, 4), (5, ), [[-2], [0]]),
((3, 5, 4), (5, ), [[1], [0]]),
((2, ), (2, 3), 1),
((3, ), (3, ), 0),
((2, ), (2, 3), 0),
((3, 5, 4), (5, ), 0),
((2, 3, 4), (4, 3, 2), [[], []]),
((3, 0), (0, 5), 1),
((3, 0), (0, 4), [[1], [0]]),
((0, 3), (3, 5), 1),
((0, 3), (5, 0), [[0], [1]])
]
for hybridize in [True, False]:
for a_shape, b_shape, axes in tensor_shapes:
for dtype in [_np.float32, _np.float64]:
test_tensordot = TestTensordot(axes)
if hybridize:
test_tensordot.hybridize()
a = rand_ndarray(shape=a_shape, dtype=dtype).as_np_ndarray()
b = rand_ndarray(shape=b_shape, dtype=dtype).as_np_ndarray()
a.attach_grad()
b.attach_grad()
np_out = _np.tensordot(a.asnumpy(), b.asnumpy(), axes)
with mx.autograd.record():
mx_out = test_tensordot(a, b)
assert mx_out.shape == np_out.shape
assert_almost_equal(mx_out.asnumpy(),
np_out,
rtol=1e-3,
atol=1e-5)
mx_out.backward()
np_backward = tensordot_backward(a.asnumpy(), b.asnumpy(),
axes)
assert_almost_equal(a.grad.asnumpy(),
np_backward[0],
rtol=1e-3,
atol=1e-5)
assert_almost_equal(b.grad.asnumpy(),
np_backward[1],
rtol=1e-3,
atol=1e-5)
# Test imperative once again
mx_out = np.tensordot(a, b, axes)
np_out = _np.tensordot(a.asnumpy(), b.asnumpy(), axes)
assert_almost_equal(mx_out.asnumpy(),
np_out,
rtol=1e-3,
atol=1e-5)
# test numeric gradient
if (_np.prod(a_shape) > 0 and _np.prod(b_shape) > 0):
a_sym = mx.sym.Variable("a").as_np_ndarray()
b_sym = mx.sym.Variable("b").as_np_ndarray()
mx_sym = mx.sym.np.tensordot(a_sym, b_sym,
axes).as_nd_ndarray()
check_numeric_gradient(
mx_sym, [a.as_nd_ndarray(),
b.as_nd_ndarray()],
rtol=1e-1,
atol=1e-1,
dtype=dtype)
def loss_plot_heatmap(physics_model, symm_type, N, max_change):
model_name = "{0}_{1}".format(physics_model, symm_type)
results_folder = "results/{0}_results".format(model_name)
study_name = "N{0}_nh{1}_lr{2}_ep{3}".format(N, NH, LR, EP)
# Define the intervals which alpha, beta will take
alpha_min, alpha_max = -max_change, max_change
beta_min, beta_max = -max_change, max_change
alpha_vals = np.linspace(alpha_min, alpha_max, DIM)
beta_vals = np.linspace(beta_min, beta_max, DIM)
grid_alphas, grid_betas = np.meshgrid(alpha_vals, beta_vals)
landscape_file = "loss_landscape_range{0}_dim{1}_{2}_{3}_seed{4}.txt".format(
max_change, DIM, model_name, study_name, SEED)
loss_data = np.loadtxt("{0}/{1}/{2}".format(results_folder, study_name,
landscape_file))
# Now make plots of parameters trajectory
# First retrieve final values of params
max_ep = 2000 # how long training occurs for
opt_params = []
for par_num in range(NUM_PARS):
param_filename = "param{0}_{1}_epoch2000_N{2}_nh{3}_lr{4}_ep{5}.txt".format(
par_num, model_name, N, NH, LR, EP)
param = np.loadtxt("{0}/{1}/param_vals/{2}".format(
results_folder, study_name, param_filename))
opt_params.append(param)
# Get recorded values from data files
period = 5 # how often parameter info recorded
epochs = range(0, max_ep + 1, period)
epochs_dict = {} # store epoch: listof params
for epoch in epochs:
params = []
for par_num in range(NUM_PARS):
param_filename = "param{0}_{1}_epoch{2}_N{3}_nh{4}_lr{5}_ep{6}.txt".format(
par_num, model_name, epoch, N, NH, LR, EP)
param = np.loadtxt("{0}/{1}/param_vals/{2}".format(
results_folder, study_name, param_filename))
param -= opt_params[par_num] # centre at optimum
params.append(param)
epochs_dict[epoch] = params
# Get deltas and etas using same seed
deltas = []
etas = []
for par_num in range(NUM_PARS):
delta = np.loadtxt(
"{0}/{1}/delta{2}_{3}_N{4}_nh{5}_lr{6}_ep{7}.txt".format(
results_folder, study_name, par_num, model_name, N, NH, LR,
EP))
eta = np.loadtxt(
"{0}/{1}/eta{2}_{3}_N{4}_nh{5}_lr{6}_ep{7}.txt".format(
results_folder, study_name, par_num, model_name, N, NH, LR,
EP))
deltas.append(delta)
etas.append(eta)
# Now compute alphas and betas for every epoch
alphas = []
betas = []
for epoch in epochs:
alpha = 0
beta = 0
for par_num in range(NUM_PARS):
# Need to expand dims to use tensor dot
curr_params = epochs_dict[epoch][par_num]
delta = deltas[par_num]
eta = etas[par_num]
if len(curr_params.shape) == 1:
curr_params = np.expand_dims(curr_params, axis=0)
delta = np.expand_dims(delta, axis=0)
eta = np.expand_dims(eta, axis=0)
alpha += np.tensordot(curr_params, delta)
beta += np.tensordot(curr_params, eta)
alphas.append(alpha)
betas.append(beta)
# NOTE: This stuff is for interpolating, which isn't good in this case
# n_points = 10000
# alpha_indices = np.random.choice(DIM, n_points)
# beta_indices = np.random.choice(DIM, n_points)
# alpha_points = alpha_vals[alpha_indices]
# beta_points = beta_vals[beta_indices]
# loss_points = loss_data[alpha_indices, beta_indices]
# points_tuple = (alpha_points, beta_points)
# grid_tuple = (grid_alphas, grid_betas)
# interp_vals = griddata(
# points_tuple, loss_points, grid_tuple, method="cubic"
# )
fig, ax = plt.subplots()
# ax.contourf(grid_alphas, grid_betas, interp_vals)
norm = colors.Normalize(vmin=loss_data.min(), vmax=loss_data.max())
contour = ax.imshow(
loss_data,
cmap="YlGn",
norm=norm,
extent=[-max_change, max_change, -max_change, max_change],
)
ax.set_xlim(-max_change, max_change)
ax.set_ylim(-max_change, max_change)
fig.colorbar(contour, ax=ax)
ax.plot(alphas, betas, color="C0", linestyle="-")
ax.scatter(alphas[::40], betas[::40], color="C3")
ax.set_title(r"{0} loss landscape projection".format(
NICE_NAMES_DICT[model_name]))
ax.set_xlabel(r"$\alpha$")
ax.set_ylabel(r"$\beta$")
# ax.set_zlabel(r"Loss")
image_name = (
"plot_loss_landscape_heatmap_range{0}_dim{1}_{2}_{3}_seed{4}.pdf".
format(max_change, DIM, model_name, study_name, SEED))
plt.savefig("{0}/{1}/{2}".format(results_folder, study_name, image_name))
def Single_Site_Rho(self, k, subspace=0):
L = self.sim['L']
da = self.sim['da']
db = self.sim['db']
Gamma_k = self.Gamma[k]
# Need to treat boundaries differently...
# See Mishmash thesis pg. 73 for formulas
if (k != 0 and k != L - 1):
Rho_L = np.tensordot(np.diag(self.Lambda[k - 1, :]),
np.conjugate(Gamma_k),
axes=(1, 1))
Rho_L = np.tensordot(Rho_L,
np.diag(self.Lambda[k, :]),
axes=(-1, 0))
Rho_R = np.tensordot(np.diag(self.Lambda[k, :]),
Gamma_k,
axes=(1, 1))
Rho_R = np.tensordot(Rho_R,
np.diag(self.Lambda[k, :]),
axes=(-1, 0))
Rho = np.tensordot(Rho_L, Rho_R, axes=([0, 2], [0, 2]))
Rho = np.transpose(Rho)
elif (k == 0):
Rho_L = np.tensordot(np.conjugate(Gamma_k),
np.diag(self.Lambda[k, :]),
axes=(-1, 0))
Rho_R = np.tensordot(Gamma_k,
np.diag(self.Lambda[k, :]),
axes=(-1, -1))
Rho = np.tensordot(Rho_L, Rho_R, axes=(-1, -1))
Rho = np.transpose(Rho)
elif (k == L - 1):
Rho_L = np.tensordot(np.diag(self.Lambda[k - 1, :]),
np.conjugate(Gamma_k),
axes=(-1, -1))
Rho_R = np.tensordot(np.diag(self.Lambda[k - 1, :]),
Gamma_k,
axes=(-1, -1))
Rho = np.tensordot(Rho_L, Rho_R, axes=(0, 0))
Rho = np.transpose(Rho)
if (subspace == 0):
return Rho
else:
Rho = np.reshape(Rho, (da, db, da, db))
if (subspace == 'a'):
Rho = np.tensordot(np.identity(db), Rho, axes=([0, 1], [1, 3]))
elif (subspace == 'b'):
Rho = np.tensordot(np.identity(da), Rho, axes=([0, 1], [0, 2]))
return Rho
def distortion(x,y,matrix):
s = np.stack(arrays=(x,y,np.ones(shape=x.shape)),axis = 0)
s = np.tensordot(matrix,s,axes=(1,0))
return s[0:2]
for x_elem, t_elem in zip(x_data, t_data):
x = Variable(np.array([x_elem], dtype=np.float32))
t = Variable(np.array([t_elem], dtype=np.int32))
y = model(x)
model.cleargrads()
loss = F.softmax_cross_entropy(y, t)
loss.backward()
# save grad
for path, param in model.namedparams():
dic_grads[path].append(param.grad)
# manipulate grad
z = np.array(z_data, dtype=np.float32)
for path, param in model.namedparams():
grads = np.array(dic_grads[path], dtype=np.float32)
grad = np.tensordot(z, grads, axes=1)
grad /= len(z_data)
param.grad = grad
optimizer.update()
# print param data and grad
for path, param in model.namedparams():
print(path)
print(param.data)
print(param.grad)
def __call__(self, M: np.array, A: np.array = None, B: np.array = None) -> (np.array, np.array):
erase = tensordot(transpose(ones_like(A) - A, axes=[1, 0]), ones_like(A), axes=1)
contrib = tensordot(transpose(A, axes=[1, 0]), B, axes=1)
new_mem = (erase * M) + contrib
return new_mem, np.array([[1 if idx == 0 else 0 for idx in np.arange(A.shape[1])]], dtype=np.float64)
def matrixRepresentation(self, vectorRepresentation):
return np.tensordot(vectorRepresentation, self.basis, axes=[0, 0])
def mat_vec_product(self, mat, vec, device_wire_labels):
r"""Apply multiplication of a matrix to subsystems of the quantum state.
Args:
mat (array): matrix to multiply
vec (array): state vector to multiply
device_wire_labels (Sequence[int]): labels of device subsystems
Returns:
array: output vector after applying ``mat`` to input ``vec`` on specified subsystems
"""
num_wires = len(device_wire_labels)
if mat.shape != (2 ** num_wires, 2 ** num_wires):
raise ValueError(
f"Please specify a {2**num_wires} x {2**num_wires} matrix for {num_wires} wires."
)
# first, we need to reshape both the matrix and vector
# into blocks of 2x2 matrices, in order to do the higher
# order matrix multiplication
# Reshape the matrix to ``size=[2, 2, 2, ..., 2]``,
# where ``len(size) == 2*len(wires)``
#
# The first half of the dimensions correspond to a
# 'ket' acting on each wire/qubit, while the second
# half of the dimensions correspond to a 'bra' acting
# on each wire/qubit.
#
# E.g., if mat = \sum_{ijkl} (c_{ijkl} |ij><kl|),
# and wires=[0, 1], then
# the reshaped dimensions of mat are such that
# mat[i, j, k, l] == c_{ijkl}.
mat = np.reshape(mat, [2] * len(device_wire_labels) * 2)
# Reshape the state vector to ``size=[2, 2, ..., 2]``,
# where ``len(size) == num_wires``.
# Each wire corresponds to a subsystem.
#
# E.g., if vec = \sum_{ijk}c_{ijk}|ijk>,
# the reshaped dimensions of vec are such that
# vec[i, j, k] == c_{ijk}.
vec = np.reshape(vec, [2] * self.num_wires)
# Calculate the axes on which the matrix multiplication
# takes place. For the state vector, this simply
# corresponds to the requested wires. For the matrix,
# it is the latter half of the dimensions (the 'bra' dimensions).
#
# For example, if num_wires=3 and wires=[2, 0], then
# axes=((2, 3), (2, 0)). This is equivalent to doing
# np.einsum("ijkl,lnk", mat, vec).
axes = (np.arange(len(device_wire_labels), 2 * len(device_wire_labels)), device_wire_labels)
# After the tensor dot operation, the resulting array
# will have shape ``size=[2, 2, ..., 2]``,
# where ``len(size) == num_wires``, corresponding
# to a valid state of the system.
tdot = np.tensordot(mat, vec, axes=axes)
# Tensordot causes the axes given in `wires` to end up in the first positions
# of the resulting tensor. This corresponds to a (partial) transpose of
# the correct output state
# We'll need to invert this permutation to put the indices in the correct place
unused_idxs = [idx for idx in range(self.num_wires) if idx not in device_wire_labels]
perm = device_wire_labels + unused_idxs
# argsort gives the inverse permutation
inv_perm = np.argsort(perm)
state_multi_index = np.transpose(tdot, inv_perm)
return np.reshape(state_multi_index, 2 ** self.num_wires)
def neighbours(df, atoms_list):
"""Generates a list of distances to neighbours (NNs) for each pairs of species: AA, AB, BA, BB."""
# The initial unit cell and supercell, separated into sublattices of A and B atoms
clA, clB = unitcell(atoms_list)
spclA, spclB = supercell(atoms_list)
# Unit vectors used to reshape the arrays
eA = np.ones(len(clA))
eB = np.ones(len(clB))
EA = np.ones(len(spclA))
EB = np.ones(len(spclB))
# Yields the matrix to calcualte the interatomic distances in the lattice coordinate system
abc = cell_parameters(df)[0]
ang = np.cos(cell_parameters(df)[1])
mat = np.array([[1, ang[2], ang[1]], [ang[2], 1, ang[0]],
[ang[1], ang[0], 1]])
mat = np.multiply(mat, np.tensordot(abc, abc, axes=0))
# Pairwise distances between atoms in the initial cell and supercell
drAA = np.tensordot(eA, spclA, axes=0) - np.swapaxes(
np.tensordot(EA, clA, axes=0), 0, 1)
drAB = np.tensordot(eA, spclB, axes=0) - np.swapaxes(
np.tensordot(EB, clA, axes=0), 0, 1)
drBA = np.tensordot(eB, spclA, axes=0) - np.swapaxes(
np.tensordot(EA, clB, axes=0), 0, 1)
drBB = np.tensordot(eB, spclB, axes=0) - np.swapaxes(
np.tensordot(EB, clB, axes=0), 0, 1)
# This part creates a 3x3 matrix for each pair to calculate the distances
drAA = np.tensordot(drAA, np.ones(3), axes=0)
drAB = np.tensordot(drAB, np.ones(3), axes=0)
drBA = np.tensordot(drBA, np.ones(3), axes=0)
drBB = np.tensordot(drBB, np.ones(3), axes=0)
drAA = np.multiply(drAA, np.swapaxes(drAA, 2, 3))
drAB = np.multiply(drAB, np.swapaxes(drAB, 2, 3))
drBA = np.multiply(drBA, np.swapaxes(drBA, 2, 3))
drBB = np.multiply(drBB, np.swapaxes(drBB, 2, 3))
# Calculate interatomic distances using the previously derived matrix
distAA = np.multiply(
np.tensordot(eA, np.tensordot(EA, mat, axes=0), axes=0), drAA)
distAB = np.multiply(
np.tensordot(eA, np.tensordot(EB, mat, axes=0), axes=0), drAB)
distBA = np.multiply(
np.tensordot(eB, np.tensordot(EA, mat, axes=0), axes=0), drBA)
distBB = np.multiply(
np.tensordot(eB, np.tensordot(EB, mat, axes=0), axes=0), drBB)
distAA = np.sqrt(np.sum(np.sum(distAA, axis=3), axis=2))
distAB = np.sqrt(np.sum(np.sum(distAB, axis=3), axis=2))
distBA = np.sqrt(np.sum(np.sum(distBA, axis=3), axis=2))
distBB = np.sqrt(np.sum(np.sum(distBB, axis=3), axis=2))
return distAA, distAB, distBA, distBB
def idwht_bloque(p): h_wh = H_WH(4) r = np.tensordot(np.tensordot(h_wh, p, axes = ([1][0])), h_wh, axes = ([1][0])) return r
def check_hop(atoms, energy_kin, energy_pot, gap, force_upper_t2,
force_upper_t1, force_lower_t2, force_lower_t1,
target_state):
global flag_es, j_md
global coordinates_t1, velocities_t1, dt, tau_0
global hop, do_hop
"""Check for local minimum and hopping attempt"""
etot = energy_pot + energy_kin
p_zn = 0.0
p_lz = 0.0
small_dgap = False
v_diab = 0.0
# check for the local minimum of the energy gap
if (gap[j_md - 1] < gap[j_md]) and (gap[j_md - 2] > gap[j_md - 1]):
# finite difference calculation of second order derivative
dgap = (gap[j_md] - 2.0 * gap[j_md - 1] + gap[j_md - 2]) / (dt *
dt)
# compute Belyaev-Lebedev probability
if (dgap < 1E-12):
#print('small or negative d^2/dt^2',dgap)
small_dgap = True
else:
c_ij = np.power(gap[j_md - 1], 3) / dgap
p_lz = np.exp(-0.5 * np.pi * np.sqrt(c_ij))
# output second time derivative between S2 and S3
#if target_state==2:
# df2.write('{0:0.2f} {1} {2} \n'.format(dt*(j_md-1)*tau_0,\
# dgap,np.power(gap[j_md-1],3)))
# compute diabatic gradients for ZN
sum_G = force_upper_t2 + force_lower_t2
dGc = (force_upper_t2 - force_lower_t2) - (force_upper_t1 -
force_lower_t1)
dGc /= dt
# dGc*velo has to be in a.u.
conversion_velo = fs * tau_0 / Bohr
dGxVelo = np.tensordot(dGc, velocities_t1 * conversion_velo)
if (dGxVelo < 0.0):
if not small_dgap:
factor = 0.5 * np.sqrt(gap[j_md - 1] / dgap)
else:
factor = 0.0
else:
factor = 0.5 * np.sqrt(gap[j_md - 1] / dGxVelo)
force1_x = 0.5 * sum_G - factor * dGc
force2_x = 0.5 * sum_G + factor * dGc
force_diff_acc = 0.0
force_prod_acc = 0.0
for i in range(0, len(a)):
for j_coord in range(0, 3):
temp0 = force2_x[i, j_coord] - force1_x[i, j_coord]
temp = temp0 * temp0
temp2 = force2_x[i, j_coord] * force1_x[i, j_coord]
force_diff_acc += temp / masses[i]
force_prod_acc += temp2 / masses[i]
force_diff = np.sqrt(force_diff_acc)
if force_prod_acc > 0.0:
do_plus = True
else:
do_plus = False
force_prod = np.sqrt(np.abs(force_prod_acc))
v_diab = gap[j_md - 1] / 2.0
a2 = 0.5 * force_diff * force_prod / (np.power(2.0 * v_diab, 3))
b2 = energy_kin * force_diff / (force_prod * 2.0 * v_diab)
if (a2 < 0.0) or (b2 < 0.0):
print('Alert!')
root = 0.0
if do_plus:
root = b2 + np.sqrt(b2 * b2 + 1.0)
else:
root = b2 + np.sqrt(np.abs(b2 * b2 - 1.0))
# compute Zhu-Nakamura probability
if not small_dgap:
p_zn = np.exp(-np.pi * 0.25 * np.sqrt(2.0 / (a2 * root)))
else:
p_zn = 0.0
# comparison with a random number
if do_hop and (not hop):
xi = np.random.rand(1)
if xi <= p_lz and do_lz:
#print('Attempted hop according to BL at {}'.format(j_md-1))
hop = True
elif xi <= p_zn and do_zn:
#print('Attempted hop according to ZN at {}'.format(j_md-1))
hop = True
betta = gap[j_md - 1] / energy_kin
# check for frustrated hop condition
if (hop and betta > 1.0 and target_state > flag_es):
hop = False
if hop:
# output energy gap at the hopping point
print('{0} {1}'.format(target_state, gap[j_md - 1]))
# make one MD step back since local minimum was at t and we are at t+dt
a.set_positions(coordinates_t1)
# velocity rescaling to conserve total energy
if target_state < flag_es:
a.set_velocities(np.sqrt(1.0 + betta) * velocities_t1)
else:
a.set_velocities(np.sqrt(1.0 - betta) * velocities_t1)
# change the running state
flag_es = target_state
if do_zn:
return p_zn
else:
return p_lz
def idwht_bloque(p):
h_wh = H_WH(4)
r = np.tensordot(np.tensordot(h_wh, p, axes = ([1][0])), h_wh, axes = ([1][0]))
return r
"""
Reproducir los bloques base de la transformación para los casos N=4,8,16
Ver imágenes adjuntas
"""
N = 4
while N < 32:
H = H_WH(N)
for row in range(N):
for col in range(N):
baseImage = np.tensordot(H[row], np.transpose(H[col]), 0)
print(baseImage)
plt.imshow(baseImage)
plt.xticks([])
plt.yticks([])
plt.show()
N *= 2
"""
MAIN
"""
Q = np.array(
[[217, 8, 248, 199],
[215, 189, 242, 10],
[200, 65, 191, 92],
[174, 239, 237, 118]]
def main(input_animation_file, output_sploc_file, output_animation_file):
rest_shape = "first" # which frame to use as rest-shape ("first" or "average")
K = 50 # number of components
smooth_min_dist = 0.1 # minimum geodesic distance for support map, d_min_in paper
smooth_max_dist = 0.7 # maximum geodesic distance for support map, d_max in paper
num_iters_max = 10 # number of iterations to run
sparsity_lambda = 2. # sparsity parameter, lambda in the paper
rho = 10.0 # penalty parameter for ADMM
num_admm_iterations = 10 # number of ADMM iterations
# preprocessing: (external script)
# rigidly align sequence
# normalize into -0.5 ... 0.5 box
with h5py.File(input_animation_file, 'r') as f:
verts = f['verts'].value.astype(np.float)
tris = f['tris'].value
F, N, _ = verts.shape
if rest_shape == "first":
Xmean = verts[0]
elif rest_shape == "average":
Xmean = np.mean(verts, axis=0)
# prepare geodesic distance computation on the restpose mesh
compute_geodesic_distance = GeodesicDistanceComputation(Xmean, tris)
# form animation matrix, subtract mean and normalize
# notice that in contrast to the paper, X is an array of shape (F, N, 3) here
X = verts - Xmean[np.newaxis]
pre_scale_factor = 1 / np.std(X)
X *= pre_scale_factor
R = X.copy() # residual
# find initial components explaining the residual
C = []
W = []
for k in range(K):
# find the vertex explaining the most variance across the residual animation
magnitude = (R**2).sum(axis=2)
idx = np.argmax(magnitude.sum(axis=0))
# find linear component explaining the motion of this vertex
U, s, Vt = la.svd(R[:,idx,:].reshape(R.shape[0], -1).T, full_matrices=False)
wk = s[0] * Vt[0,:] # weights
# invert weight according to their projection onto the constraint set
# this fixes problems with negative weights and non-negativity constraints
wk_proj = project_weight(wk)
wk_proj_negative = project_weight(-wk)
wk = wk_proj \
if norm(wk_proj) > norm(wk_proj_negative) \
else wk_proj_negative
s = 1 - compute_support_map(idx, compute_geodesic_distance, smooth_min_dist, smooth_max_dist)
# solve for optimal component inside support map
ck = (np.tensordot(wk, R, (0, 0)) * s[:,np.newaxis])\
/ np.inner(wk, wk)
C.append(ck)
W.append(wk)
# update residual
R -= np.outer(wk, ck).reshape(R.shape)
C = np.array(C)
W = np.array(W).T
# prepare auxiluary variables
Lambda = np.empty((K, N))
U = np.zeros((K, N, 3))
# main global optimization
for it in range(num_iters_max):
# update weights
Rflat = R.reshape(F, N*3) # flattened residual
for k in range(C.shape[0]): # for each component
Ck = C[k].ravel()
Ck_norm = np.inner(Ck, Ck)
if Ck_norm <= 1.e-8:
# the component seems to be zero everywhere, so set it's activation to 0 also
W[:,k] = 0
continue # prevent divide by zero
# block coordinate descent update
Rflat += np.outer(W[:,k], Ck)
opt = np.dot(Rflat, Ck) / Ck_norm
W[:,k] = project_weight(opt)
Rflat -= np.outer(W[:,k], Ck)
# update spatially varying regularization strength
for k in range(K):
ck = C[k]
# find vertex with biggest displacement in component and compute support map around it
idx = (ck**2).sum(axis=1).argmax()
support_map = compute_support_map(idx, compute_geodesic_distance,
smooth_min_dist, smooth_max_dist)
# update L1 regularization strength according to this support map
Lambda[k] = sparsity_lambda * support_map
# update components
Z = C.copy() # dual variable
# prefactor linear solve in ADMM
G = np.dot(W.T, W)
c = np.dot(W.T, X.reshape(X.shape[0], -1))
solve_prefactored = cho_factor(G + rho * np.eye(G.shape[0]))
# ADMM iterations
for admm_it in range(num_admm_iterations):
C = cho_solve(solve_prefactored, c + rho * (Z - U).reshape(c.shape)).reshape(C.shape)
Z = prox_l1l2(Lambda, C + U, 1. / rho)
U = U + C - Z
# set updated components to dual Z,
# this was also suggested in [Boyd et al.] for optimization of sparsity-inducing norms
C = Z
# evaluate objective function
R = X - np.tensordot(W, C, (1, 0)) # residual
sparsity = np.sum(Lambda * np.sqrt((C**2).sum(axis=2)))
e = (R**2).sum() + sparsity
# TODO convergence check
print("iteration %03d, E=%f" % (it, e))
# undo scaling
C /= pre_scale_factor
# save components
with h5py.File(output_sploc_file, 'w') as f:
f['default'] = Xmean
f['tris'] = tris
for i, c in enumerate(C):
f['comp%03d' % i] = c + Xmean
print(np.shape(W))
# save encoded animation including the weights
if output_animation_file:
with h5py.File(output_animation_file, 'w') as f:
f['verts'] = np.tensordot(W, C, (1, 0)) + Xmean[np.newaxis]
f['tris'] = tris
f['weights'] = W
